Add prompt to code service

Author: intitni

Core/Package.swift

@@ -87,6 +87,7 @@ let package = Package(
                 "AXExtension",
                 "Logger",
                 "ChatService",
+                "PromptToCodeService",
                 .product(name: "AsyncAlgorithms", package: "swift-async-algorithms"),
             ]
         ),
@@ -154,5 +155,10 @@ let package = Package(
         .target(name: "ChatPlugins", dependencies: ["OpenAIService", "Environment", "Terminal"]),
         .target(name: "Terminal"),
         .target(name: "ChatService", dependencies: ["OpenAIService", "ChatPlugins", "Environment"]),
+        .target(
+            name: "PromptToCodeService",
+            dependencies: ["OpenAIService", "Environment", "CopilotService", "CopilotModel"]
+        ),
+        .testTarget(name: "PromptToCodeServiceTests", dependencies: ["PromptToCodeService"]),
     ]
 )

Core/Sources/PromptToCodeService/PromptToCodeService.swift

@@ -0,0 +1,220 @@
+import CopilotModel
+import CopilotService
+import Foundation
+import OpenAIService
+
+public final class PromptToCodeService: ObservableObject {
+    var designatedPromptToCodeAPI: PromptToCodeAPI?
+    var promptToCodeAPI: PromptToCodeAPI {
+        designatedPromptToCodeAPI ?? OpenAIPromptToCodeAPI()
+    }
+    var runningAPI: PromptToCodeAPI?
+
+    @Published public var oldCode: String?
+    @Published public var code: String
+    @Published public var isResponding: Bool = false
+    @Published public var description: String = ""
+    public var canRevert: Bool { oldCode != nil }
+    public var selectionRange: CursorRange
+    public var language: CopilotLanguage
+    public var indentSize: Int
+    public var usesTabsForIndentation: Bool
+
+    public init(
+        code: String,
+        selectionRange: CursorRange,
+        language: CopilotLanguage,
+        identSize: Int,
+        usesTabsForIndentation: Bool
+    ) {
+        self.code = code
+        self.selectionRange = selectionRange
+        self.language = language
+        indentSize = identSize
+        self.usesTabsForIndentation = usesTabsForIndentation
+    }
+
+    public func modifyCode(prompt: String) async throws {
+        let api = promptToCodeAPI
+        runningAPI = api
+        isResponding = true
+        let toBemodified = code
+        oldCode = code
+        code = ""
+        description = ""
+        defer { isResponding = false }
+        let stream = try await api.modifyCode(
+            code: toBemodified,
+            language: language,
+            indentSize: indentSize,
+            usesTabsForIndentation: usesTabsForIndentation,
+            requirement: prompt
+        )
+        for try await fragment in stream {
+            Task { @MainActor in
+                code = fragment.code
+                description = fragment.description
+            }
+        }
+    }
+
+    public func revert() {
+        guard let oldCode = oldCode else { return }
+        code = oldCode
+        self.oldCode = nil
+    }
+    
+    public func generateCompletion() -> CopilotCompletion {
+        .init(
+            text: code,
+            position: selectionRange.start,
+            uuid: UUID().uuidString,
+            range: selectionRange,
+            displayText: code
+        )
+    }
+    
+    public func stopResponding() {
+        runningAPI?.stopResponding()
+        isResponding = false
+    }
+}
+
+protocol PromptToCodeAPI {
+    func modifyCode(
+        code: String,
+        language: CopilotLanguage,
+        indentSize: Int,
+        usesTabsForIndentation: Bool,
+        requirement: String
+    ) async throws -> AsyncThrowingStream<(code: String, description: String), Error>
+    
+    func stopResponding()
+}
+
+final class OpenAIPromptToCodeAPI: PromptToCodeAPI {
+    var service: (any ChatGPTServiceType)?
+    
+    func stopResponding() {
+        Task {
+            await service?.stopReceivingMessage()
+        }
+    }
+    
+    func modifyCode(
+        code: String,
+        language: CopilotLanguage,
+        indentSize: Int,
+        usesTabsForIndentation: Bool,
+        requirement: String
+    ) async throws -> AsyncThrowingStream<(code: String, description: String), Error> {
+        let prompt = {
+            let indentRule = usesTabsForIndentation ? "\(indentSize) tabs" : "\(indentSize) spaces"
+            if code.isEmpty {
+                return """
+                You are a senior programer in writing code in \(language.rawValue).
+
+                Please write a piece of code that meets my requirements. The indentation should be \(
+                    indentRule
+                ).
+
+                Please reply to me start with the code block, followed by a short description in 1-3 sentences about what you did.
+                """
+            } else {
+                return """
+                You are a senior programer in writing code in \(language.rawValue).
+
+                Please mutate the following code with my requirements. The indentation should be \(
+                    indentRule
+                ).
+
+                Please reply to me start with the code block followed by a short description about what you did in 1-3 sentences.
+
+                ```
+                \(code)
+                ```
+                """
+            }
+        }()
+
+        let chatGPTService = ChatGPTService(systemPrompt: prompt)
+        service = chatGPTService
+        let stream = try await chatGPTService.send(content: requirement)
+        return .init { continuation in
+            Task {
+                var content = ""
+                do {
+                    for try await fragment in stream {
+                        content.append(fragment)
+                        continuation.yield(extractCodeAndDescription(from: content))
+                    }
+                    continuation.finish()
+                } catch {
+                    continuation.finish(throwing: error)
+                }
+            }
+        }
+    }
+
+    func extractCodeAndDescription(from content: String) -> (code: String, description: String) {
+        func extractCodeFromMarkdown(_ markdown: String) -> (code: String, endIndex: Int)? {
+            let codeBlockRegex = try! NSRegularExpression(
+                pattern: #"```(?:\w+)?[\n]([\s\S]+?)[\n]```"#,
+                options: .dotMatchesLineSeparators
+            )
+            let range = NSRange(markdown.startIndex..<markdown.endIndex, in: markdown)
+            if let match = codeBlockRegex.firstMatch(in: markdown, options: [], range: range) {
+                let codeBlockRange = Range(match.range(at: 1), in: markdown)!
+                return (String(markdown[codeBlockRange]), match.range(at: 0).upperBound)
+            }
+
+            let incompleteCodeBlockRegex = try! NSRegularExpression(
+                pattern: #"```(?:\w+)?[\n]([\s\S]+?)$"#,
+                options: .dotMatchesLineSeparators
+            )
+            let range2 = NSRange(markdown.startIndex..<markdown.endIndex, in: markdown)
+            if let match = incompleteCodeBlockRegex.firstMatch(
+                in: markdown,
+                options: [],
+                range: range2
+            ) {
+                let codeBlockRange = Range(match.range(at: 1), in: markdown)!
+                return (String(markdown[codeBlockRange]), match.range(at: 0).upperBound)
+            }
+            return nil
+        }
+
+        guard let (code, endIndex) = extractCodeFromMarkdown(content) else {
+            return ("", "")
+        }
+
+        func extractDescriptionFromMarkdown(_ markdown: String, startIndex: Int) -> String {
+            let startIndex = markdown.index(markdown.startIndex, offsetBy: startIndex)
+            guard startIndex < markdown.endIndex else { return "" }
+            let range = startIndex..<markdown.endIndex
+            let description = String(markdown[range])
+                .trimmingCharacters(in: .whitespacesAndNewlines)
+            return description
+        }
+
+        let description = extractDescriptionFromMarkdown(content, startIndex: endIndex)
+        
+        return (code, description)
+    }
+}
+
+final class CopilotPromptToCodeAPI: PromptToCodeAPI {
+    func stopResponding() {
+        fatalError()
+    }
+    
+    func modifyCode(
+        code: String,
+        language: CopilotLanguage,
+        indentSize: Int,
+        usesTabsForIndentation: Bool,
+        requirement: String
+    ) async throws -> AsyncThrowingStream<(code: String, description: String), Error> {
+        fatalError()
+    }
+}

Core/Tests/PromptToCodeServiceTests/ExtractCodeFromChatGPTTests.swift

@@ -0,0 +1,69 @@
+import XCTest
+@testable import PromptToCodeService
+
+final class ExtractCodeFromChatGPTTests: XCTestCase {
+    func test_extract_from_no_code_block() {
+        let api = OpenAIPromptToCodeAPI()
+        let result = api.extractCodeAndDescription(from: """
+        hello world!
+        """)
+        
+        XCTAssertEqual(result.code, "")
+        XCTAssertEqual(result.description, "")
+    }
+    
+    func test_extract_from_incomplete_code_block() {
+        let api = OpenAIPromptToCodeAPI()
+        let result = api.extractCodeAndDescription(from: """
+        ```swift
+        func foo() {}
+        """)
+        
+        XCTAssertEqual(result.code, "func foo() {}")
+        XCTAssertEqual(result.description, "")
+    }
+    
+    func test_extract_from_complete_code_block() {
+        let api = OpenAIPromptToCodeAPI()
+        let result = api.extractCodeAndDescription(from: """
+        ```swift
+        func foo() {}
+        
+        func bar() {}
+        ```
+        
+        Description
+        """)
+        
+        XCTAssertEqual(result.code, "func foo() {}\n\nfunc bar() {}")
+        XCTAssertEqual(result.description, "Description")
+    }
+    
+    func test_extract_from_incomplete_code_block_without_language() {
+        let api = OpenAIPromptToCodeAPI()
+        let result = api.extractCodeAndDescription(from: """
+        ```
+        func foo() {}
+        """)
+        
+        XCTAssertEqual(result.code, "func foo() {}")
+        XCTAssertEqual(result.description, "")
+    }
+    
+    func test_extract_from_code_block_without_language() {
+        let api = OpenAIPromptToCodeAPI()
+        let result = api.extractCodeAndDescription(from: """
+        ```
+        func foo() {}
+        
+        func bar() {}
+        ```
+        
+        Description
+        """)
+        
+        XCTAssertEqual(result.code, "func foo() {}\n\nfunc bar() {}")
+        XCTAssertEqual(result.description, "Description")
+    }
+    
+}

TestPlan.xctestplan

@@ -63,6 +63,13 @@
         "identifier" : "SuggestionWidgetTests",
         "name" : "SuggestionWidgetTests"
       }
+    },
+    {
+      "target" : {
+        "containerPath" : "container:Core",
+        "identifier" : "PromptToCodeServiceTests",
+        "name" : "PromptToCodeServiceTests"
+      }
     }
   ],
   "version" : 1