RefactorGPT 😾

README.md

@@ -2,7 +2,9 @@
 
 A Zsh helper to auto-correct the current command via ChatGPT.
 
-Just compile the go code to executable. Code is written _badly_ but it does its job:
+Just compile the go code to executable.
+
+~~Code is written _badly_ but it does its job~~ Code is now edited with ChatGPT, since it's better than me in golang 💩
 
 - Looks for an API key in CHATGPT_API_KEY environment variable
 - It accepts one argument, intended to get a whole command inside this argument

oh-my-gpt-correction.go

@@ -12,17 +12,49 @@ import (
 )
 
 func main() {
-	apiKey := os.Getenv("CHATGPT_API_KEY")
+	// Validate command-line arguments
+	validateArguments()
 
+	// Get API key and query from environment variable and command-line argument
+	apiKey := os.Getenv("CHATGPT_API_KEY")
 	query := os.Args[1]
-	if query == "" {
+
+	// Create OpenAI client
+	client := openai.NewClient(apiKey)
+
+	// Request chat completion from the API
+	result, err := requestChatCompletion(client, query)
+
+	if err != nil {
+		fmt.Printf("ChatCompletion error: %v\n", err)
+		os.Exit(1)
+	}
+
+	// Clean up the result by removing unnecessary whitespace characters
+	trimmed := cleanResult(result)
+
+	// If the result is "ZOMK", exit with status code 3
+	if trimmed == "ZOMK" {
+		os.Exit(3)
+	}
+
+	// Print the comparison between the original query and the result
+	printComparison(query, trimmed)
+}
+
+// validateArguments checks if the command-line arguments and environment variable are valid
+func validateArguments() {
+	if len(os.Args) < 2 || os.Args[1] == "" {
 		os.Exit(4)
 	}
 
-	if apiKey == "" {
+	if os.Getenv("CHATGPT_API_KEY") == "" {
 		os.Exit(2)
 	}
-	client := openai.NewClient(apiKey)
+}
+
+// requestChatCompletion sends a chat completion request to the API and returns the result
+func requestChatCompletion(client *openai.Client, query string) (string, error) {
 	resp, err := client.CreateChatCompletion(
 		context.Background(),
 		openai.ChatCompletionRequest{
@@ -39,22 +71,22 @@ func main() {
 	)
 
 	if err != nil {
-		fmt.Printf("ChatCompletion error: %v\n", err)
-		os.Exit(1)
+		return "", err
 	}
 
-	result := resp.Choices[0].Message.Content
-	// No newlines pls
-	trimmed := regexp.MustCompile(`[\t\r\n]+`).ReplaceAllString(strings.TrimSpace(result), "\n")
-	if trimmed == "ZOMK" {
-		// This case is when it could not suggest anything better
-		os.Exit(3)
-	}
-
-	// Compare query with result
-	dmp := diffmatchpatch.New()
-	diffs := dmp.DiffMain(query, trimmed, false)
-	// Print both lines
-	fmt.Println(dmp.DiffPrettyText(diffs))
-	fmt.Println(trimmed)
+	return resp.Choices[0].Message.Content, nil
+}
+
+// cleanResult removes unnecessary whitespace characters from the result
+func cleanResult(result string) string {
+	return regexp.MustCompile(`[\t\r\n]+`).ReplaceAllString(strings.TrimSpace(result), "\n")
+}
+
+// printComparison prints the difference between the original query and the result
+func printComparison(query, result string) {
+	dmp := diffmatchpatch.New()
+	diffs := dmp.DiffMain(query, result, false)
+
+	fmt.Println(dmp.DiffPrettyText(diffs))
+	fmt.Println(result)
 }