Using ChatGPT to Optimize Code in Xcode: A New Era of Development

As a developer, you know that writing efficient and effective code is essential for creating successful software. However, optimizing code can be a time-consuming and challenging task. Fortunately, you can now use an Xcode extension with ChatGPT to optimize your code quickly and easily. In this article, we'll show you how to create an Xcode Source Editor Extension that connects to ChatGPT that helps you optimize your code in no time.

First, as we all currently know ChatGPT is an artificial intelligence chatbot that is built on top of OpenAI's GPT-3.5 and GPT-4 families of Large Language Models (LLMs) and fine-tuned to be applied using natural language requests.

We have to keep in mind ChatGPT will not be a replacement for developers, it will help us with a tool to make us better and improve developer quality with less effort, the human ambiguity will never be replicated to follow by an AI.

So now we are going to approach the capacity of this AI to create an Xcode Source Editor extension to use in a way that can help us to optimize our code in the fastest possible way, keep in mind that the solution provided by ChatGPT should not be considered as a final solution. Taking these considerations we'll jump right for the creation of the extension.

How to include ChatGPT to Xcode

The first thing to do is to register to OpenAI after that at the date of this publication OpenAI gives you $5.00 that you'll need to make requests to the API, after you created your account proceed to your profile to generate an API Key.

Generate your OpenAI API Key

1. Go to your profile screen located at the upper right corner and select "View API Keys"
2. Then select "Create new secret key"
3. Keep the generated key in a note because we are going to need it in a few moments to communicate to OpenAI REST API

How to create an Xcode Source Editor Extension

Now that we have everything ready for our code optimization extension it's time to proceed with the elaboration of our ChatGPT Extension. Since the release of Xcode 13 the way to create extensions has changed a bit.

1. First we need to create a new macOS app inside Xcode, you can select SwiftUI or UIKit in case you want to create a UI for your extension in this case we are going to focus only on the extension. You can save it wherever you want.
2. Go to your xcodeproj file (the first one that appears with your project name and blue square icon) to add a new target, select Xcode Source Editor Extension to embed to the project, and activate.
3. From here we are ready to start our ChatGPT extension for Xcode

Every time you want to try your extension make sure you select the correct scheme in this occasion it's called ChatGPTExtension, and select Xcode, then you can proceed to test your extension from the Editor tab and at the bottom of the list.

Coding your ChatGPT Extension

To start we'll need to create the class and codable struct to make the calls to the ChatGPT API, you could use my example or elaborate on your classes is up to you. Remember this code must select the extension target and not the main app.

• Create a new group for an OpenAI-related name it OpenAIClient, from here create a group Model to keep a better structure.

• The first struct we need is going to be called Message, it will be Codable because we require it for both cases to request and service response.
  

  struct Message: Codable {
      let role: String
      let content: String
  }
  

• Then we proceed with the structure for our ChatCompletionRequest, which is the one that will send the code that we want to optimize, it's encodable because it only sends data and doesn't expect to receive, this way we keep our code cleaner.
  

  struct ChatCompletionRequest: Encodable {
      let model: String
      let messages: [Message]
      let maxTokens: Int
  
      private enum CodingKeys: String, CodingKey {
          case model
          case messages
          case maxTokens = "max_tokens"
      }
  }
  
  

• Now it's time for the response of the services for this will create the ChatCompletionResponse that will be responsible to receive the code suggestions to optimize our code with the explanation of what was optimized from our code.
  

  struct ChatCompletionResponse: Decodable {
      let id: String
      let object: String
      let created: Int
      let model: String
      let usage: Usage
      let choices: [Choice]
  }
  
  struct Usage: Decodable {
      let promptTokens: Int
      let completionTokens: Int
      let totalTokens: Int
  
      private enum CodingKeys: String, CodingKey {
          case promptTokens = "prompt_tokens"
          case completionTokens = "completion_tokens"
          case totalTokens = "total_tokens"
      }
  }
  
  struct Choice: Decodable {
      let message: Message
      let finishReason: String
      let index: Int
  
      private enum CodingKeys: String, CodingKey {
          case message
          case finishReason = "finish_reason"
          case index
      }
  }
  
  

These structs are taken from the request and response body to consume the OpenAI Chat Completion API directly.

• It's time to proceed with OpenAI class to use our recently created models, to start creating a new swift file name OpenAI then import Foundation library we need for URLSessions calls to the API

Import Foundation

• Then for testing purposes declare two variables endpoint and apikey, at the end we are going to secure this data so it will be obtained from the plist file of our project
  

  let endpoint = "https://api.openai.com/v1/chat/completions"
  let apiKey = "<Insert your secret api key here>"
  

• Proceed with the function that creates the request to the RestAPI, it will require the prompt we introduce as a parameter it has to be a Message object that will accept and return an URLRequest, you can call it private because it only will be accessed through the class.
  

  private func createRequest(prompt: Message) -> URLRequest {    
  }
  

• Add the URLRequest with the required headers in this case the endpoint and apikey variables and a POST method, then with the prompt create a ChatCompletionRequest that will be encoded as JSON and return the created request.
  In the model you can change it to gpt-4 if you have access to it.
  

  private func createRequest(prompt: Message) -> URLRequest {
      var request = URLRequest(url: URL(string: endpoint)!)
      request.httpMethod = "POST"
      request.setValue("application/json", forHTTPHeaderField: "Content-Type")
      request.setValue("Bearer \(apiKey)", forHTTPHeaderField: "Authorization")
  
      let requestBody = ChatCompletionRequest(model: "gpt-3.5-turbo",
                                       messages: [prompt],
                                       maxTokens: 1000)
  
      let encoder = JSONEncoder()
      do {
          let jsonData = try encoder.encode(requestBody)
          request.httpBody = jsonData
      } catch {
         print(error)
      }
  
      return request
  }
  

• For the sendRequest function we are going to expect an @escaping closure to handle the response of the API, make a call to createReqeuest function and then send all the data to an URLSession with dataTask, after that send the obtained data through the closure.
  

  private func sendRequest(prompt: Message, completion: @escaping (ChatCompletionResponse?, Error?) -> ()) {
      let request = createRequest(prompt: prompt)
      let session = URLSession.shared
  
      let task = session.dataTask(with: request) { (data, response, error) in
          if let error = error {
              completion(nil, error)
              return
          }
  
          guard let httpResponse = response as? HTTPURLResponse else {
              completion(nil, NSError(domain: "OpenAIClientError", code: 0, userInfo: nil))
              return
          }
  
          guard (200...299).contains(httpResponse.statusCode) else {
              completion(nil, NSError(domain: "OpenAIClientError", code: httpResponse.statusCode, userInfo: nil))
              return
          }
  
          guard let data = data else {
              completion(nil, NSError(domain: "OpenAIClientError", code: 0, userInfo: nil))
              return
          }
  
          do {
              let decoder = JSONDecoder()
              let response = try decoder.decode(ChatCompletionResponse.self, from: data)
              completion(response, nil)
          } catch {
              completion(nil, error)
          }
      }
  
      task.resume()
  }
  

• Now will create the public function to send the prompt for the code to optimize, the way to send the expected prompt context is as easy as only asking it to optimize your code, on the demo I'll show what's the output.
  

    func getOptimizedCode(for prompt: String, completionHandler: @escaping (String?) -> Void){
        sendRequest(
            prompt: Message(role: "user", content:"""
    Optimize the following Swift code:
    ```

swift
    \(prompt)


    ```
    """)

        ) { (response, error) in
            guard error == nil else {
                print("There was an error in the OpenAI call.")
                print(error?.localizedDescription ?? "")
                completionHandler(nil)
                return
            }

            if let optimizedCode = response?.choices.first?.message.content{
                let result = "/**\n\(optimizedCode)\n*/"
                completionHandler(result)
            }
        }
    }

• And that's how the OpenAI class consists now we are going to proceed to create the command for the extension

Xcode Source Editor Command

A command in an Xcode Extension is the way to call the task or function to be executed in our source code, in this case, it will be to optimize our code.
We can have as many commands as we wish to have there are no limits, and every Command must be in its specific class that includes the protocol NSObject and XCSourceEditorCommand

• First rename the default class SourceEditorCommand to OptimizeCodeCommand from here will going to pick the selected code that we want to optimize, I chose to only selected code to not waste tokens that only have a variable in the class, also in this way we are sure what code is being optimized.

• On the Command class we have a perform function that has an invocation variable that's an XCSourceEditorCommandInvocation this object has a lot of functions but the one we are going to use is the buffer that helps us to get the text from the source, is not as easy to just select and it will appear we have to process the selected lines with the next function.
  

  private func extractSelection(_ selections: [XCSourceTextRange], fromText lines: [String]) -> String {
      return selections.compactMap { range in
          (range.start.line...range.end.line).map { lines[$0] }.joined()
      }.joined()
  }
  

• Now we are ready to process the selected lines of our source code to send it directly to ChatGPT and print the recommended code optimizations for our source code. Inside of perform function add the next code.
  

  func perform(with invocation: XCSourceEditorCommandInvocation, completionHandler: @escaping (Error?) -> Void ) -> Void {
      guard let selectedTextRanges = invocation.buffer.selections as? [XCSourceTextRange],
            let lines = invocation.buffer.lines as? [String]
      else { return completionHandler(nil) }
  
      let selectedText = extractSelection(selectedTextRanges, fromText: lines)
      OpenAIClient().getOptimizedCode(for: selectedText) { result in
          invocation.buffer.lines.add(result ?? "")
          completionHandler(nil)
      }
  }
  

The completion closure helps us to wait for the ChatGPT function to end before doing anything else, if you put the completion outside the getOptimizedCode function the command will end up as fast as we click on it.

• We have all the required code now we need to update the extension plist info to assign this command to our extension, at first we require to open the info.plist inside of the extension folder, inside of NSExtensionAttributes then XCSourceEditorCommandDefinitions every item inside of those definitions is where we can add each command of our extension.
• Change value of XCSourceEditorCommandClassName from $(PRODUCT_MODULE_NAME).SourceEditorCommand to $(PRODUCT_MODULE_NAME).<Command Class Name> eg. $(PRODUCT_MODULE_NAME).OptimizeCodeCommand
• Change value of XCSourceEditorCommandIdentifier from $(PRODUCT_BUNDLE_IDENTIFIER).SourceEditorCommand to $(PRODUCT_BUNDLE_IDENTIFIER).<Command Class Name> eg. $(PRODUCT_BUNDLE_IDENTIFIER).OptimizeCodeCommand
• Change value of XCSourceEditorCommandName to whatever you want to call your command in this case we are going to name it Optimize code with chatGPT

Now the extension is ready to be executed and try it with whatever code you want to optimize

Demo

After this preparation now we are ready to run some examples:

In this example we have a spaghetti code:

    func calculateTotalPrice(quantity: Int, price: Double, discount: Double) -> Double {
        var finalPrice = 0.0
        if quantity > 0 {
            finalPrice = price * Double(quantity)
            if discount > 0 {
                let discountAmount = finalPrice * discount / 100.0
                finalPrice -= discountAmount
                if finalPrice < 0 {
                    finalPrice = 0
                }
            }
        }
        return finalPrice
    }

The problem with this code is the amount of nested if's that are complicated that adds complexity to the code.

Now we proceed to run our app extension, I've created a new swift file, to add this function, after we select the code we go to the editor tab and select our extension name.

After we select our code to be optimized we select the command to run.

And after the process is done it will appear at the end of the file you're editing.

Protecting your API Key and endpoint

A really good way to protect our apikey or any sensitive information we can use the internal plist info to add as a key-value inside the dictionary after that we get like this way.

private var apiKey: String {
  get {
    // 1
    guard let filePath = Bundle.main.path(forResource: "Info", ofType: "plist") else {
      fatalError("Couldn't find file 'Info.plist'.")
    }
    // 2
    let plist = NSDictionary(contentsOfFile: filePath)
    guard let value = plist?.object(forKey: "API_KEY") as? String else {
      fatalError("Couldn't find key 'API_KEY' in 'Info.plist'.")
    }
    return value
  }
}

Conclusion

ChatGPT is a great tool not only for any general use, at this moment and also at any moment in the future I don't think it will take away the work for developers, but it can be used as a power-up because in this case can be used for code optimization, it still requires further analysis and testing to check that the suggested code works correctly and also if the propose code optimization works, AI will be a helper tool but not an unassisted definitive work machine, it still needs to be supervised, don't depend entirely of it but embrace it as a new tool of work.

Useful Links

Creating an Xcode Source Editor Extension

OpenAI Platform

How ChatGPT works?