Add initial refactor code

Author: chand1012

Justfile

@@ -28,6 +28,12 @@ install: build
   rm -rf $GOPATH/bin/otto
   cp bin/otto $GOPATH/bin
 
+add command:
+  cobra-cli add {{command}}
+
+test:
+  go test -v ./...
+
 crossbuild:
   #!/bin/bash
 

cmd/ask.go

@@ -27,6 +27,7 @@ var askCmd = &cobra.Command{
 	Short: "Ask a question about a file or repo",
 	Long: `Uses full text search to find relevant code and ask questions about said code.
 Requires a path to a repository or file as a positional argument.`,
+	Aliases: []string{"a"},
 	PreRun: func(cmd *cobra.Command, args []string) {
 		if verbose {
 			log.SetLevel(l.DebugLevel)

cmd/commit.go

@@ -26,9 +26,10 @@ type fileDiff struct {
 
 // commitCmd represents the commit command
 var commitCmd = &cobra.Command{
-	Use:   "commit",
-	Short: "Generates a commit message from the git diff",
-	Long:  `Uses the git diff to generate a commit message. Requires Git to be installed on the system.`,
+	Use:     "commit",
+	Short:   "Generates a commit message from the git diff",
+	Long:    `Uses the git diff to generate a commit message. Requires Git to be installed on the system.`,
+	Aliases: []string{"cm"},
 	PreRun: func(cmd *cobra.Command, args []string) {
 		if verbose {
 			log.SetLevel(l.DebugLevel)

cmd/config.go

@@ -29,6 +29,7 @@ GitHub Tokens need access to the repo scope.
 OpenAI API Key Generation: https://platform.openai.com/account/api-keys
 GitHub Token Generation: https://github.com/settings/tokens
 `,
+	Aliases: []string{"c"},
 	Run: func(cmd *cobra.Command, args []string) {
 		// load the config
 		c, err := config.Load()

cmd/docs.go

@@ -26,6 +26,7 @@ search for files in the directory and document them. If a single file is specifi
 Example:
 otto docs . -i -w 
 	`,
+	Aliases: []string{"d"},
 	PreRun: func(cmd *cobra.Command, args []string) {
 		if verbose {
 			log.SetLevel(l.DebugLevel)

cmd/refactor.go

@@ -0,0 +1,138 @@
+/*
+Copyright © 2023 NAME HERE <EMAIL ADDRESS>
+*/
+package cmd
+
+import (
+	"fmt"
+	"os"
+	"strings"
+
+	"github.com/chand1012/ottodocs/pkg/ai"
+	"github.com/chand1012/ottodocs/pkg/config"
+	"github.com/chand1012/ottodocs/pkg/constants"
+	"github.com/chand1012/ottodocs/pkg/textfile"
+	"github.com/chand1012/ottodocs/pkg/utils"
+	l "github.com/charmbracelet/log"
+	"github.com/spf13/cobra"
+)
+
+// refactorCmd represents the refactor command
+var refactorCmd = &cobra.Command{
+	Use:   "refactor",
+	Short: "A brief description of your command",
+	Long: `A longer description that spans multiple lines and likely contains examples
+and usage of using your command. For example:
+
+Cobra is a CLI library for Go that empowers applications.
+This application is a tool to generate the needed files
+to quickly create a Cobra application.`,
+	Aliases: []string{"r"},
+	PreRun: func(cmd *cobra.Command, args []string) {
+		if verbose {
+			log.SetLevel(l.DebugLevel)
+		}
+	},
+	Run: func(cmd *cobra.Command, args []string) {
+		if len(args) == 0 {
+			log.Error("Requires a file name as an argument. Example: otto refactor main.go")
+			os.Exit(1)
+		}
+
+		c, err := config.Load()
+		if err != nil {
+			log.Errorf("Error loading config: %s", err)
+			os.Exit(1)
+		}
+
+		fileName := args[0]
+
+		// load the file
+		contents, err := utils.LoadFile(fileName)
+		if err != nil {
+			log.Errorf("Error loading file: %s", err)
+			os.Exit(1)
+		}
+
+		if endLine < 0 || startLine < 0 {
+			log.Error("End line must be greater than or equal to start line and both must be greater than or equal to 0")
+			os.Exit(1)
+		}
+
+		if endLine != 0 {
+			// get the lines to refactor
+			lines := strings.Split(contents, "\n")
+			if endLine > len(lines) {
+				log.Error("End line is greater than the number of lines in the file")
+				os.Exit(1)
+			}
+
+			contents = strings.Join(lines[startLine-1:endLine], "\n")
+		} else {
+			endLine = len(strings.Split(contents, "\n"))
+		}
+
+		log.Debugf("Refactoring lines %d-%d", startLine, endLine)
+
+		if chatPrompt == "" {
+			chatPrompt, err = utils.Input("Goal: ")
+			if err != nil {
+				log.Errorf("Error prompting for goal: %s", err)
+				os.Exit(1)
+			}
+		}
+
+		prompt := constants.REFACTOR_CODE_PROMPT + "Goal: " + chatPrompt + "\n\n" + strings.TrimRight(contents, " \n")
+
+		stream, err := ai.SimpleStreamRequest(prompt, c)
+		if err != nil {
+			log.Errorf("Error requesting from OpenAI: %s", err)
+			os.Exit(1)
+		}
+
+		// print the response
+		fmt.Println("New Code:")
+		newCode, err := utils.PrintChatCompletionStream(stream)
+		if err != nil {
+			log.Errorf("Error printing chat completion stream: %s", err)
+			os.Exit(1)
+		}
+
+		if !force {
+			confirm, err := utils.Input("Would you like to overwrite the file with the new code? (y/N): ")
+			if err != nil {
+				log.Errorf("Error getting input: %s", err)
+				os.Exit(1)
+			}
+
+			confirm = strings.ToLower(confirm)
+			if confirm != "y" && confirm != "yes" {
+				os.Exit(0)
+			}
+		}
+
+		// write the new code to the file
+		finalCode, err := textfile.ReplaceLines(contents, startLine, endLine, newCode)
+		if err != nil {
+			log.Errorf("Error replacing lines: %s", err)
+			os.Exit(1)
+		}
+
+		err = utils.WriteFile(fileName, finalCode)
+		if err != nil {
+			log.Errorf("Error writing file: %s", err)
+			os.Exit(1)
+		}
+		fmt.Println("File written successfully!")
+	},
+}
+
+func init() {
+	RootCmd.AddCommand(refactorCmd)
+
+	refactorCmd.Flags().BoolVarP(&force, "force", "f", false, "Force overwrite of existing files")
+	refactorCmd.Flags().BoolVarP(&verbose, "verbose", "v", false, "Verbose output")
+	refactorCmd.Flags().IntVarP(&startLine, "start", "s", 1, "Start line")
+	refactorCmd.Flags().IntVarP(&endLine, "end", "e", 0, "End line")
+	refactorCmd.Flags().StringVarP(&chatPrompt, "goal", "g", "", "Goal of the refactor")
+}

cmd/vars.go

@@ -40,6 +40,9 @@ var useComments bool
 var promptOnly bool
 var countFinalTokens bool
 
+var startLine int
+var endLine int
+
 var log = l.NewWithOptions(os.Stderr, l.Options{
 	Level:           l.InfoLevel,
 	ReportTimestamp: false,

pkg/constants/prompts.go

@@ -39,3 +39,5 @@ var PR_TITLE_PROMPT string = "You are a helpful assistant who writes pull reques
 var PR_BODY_PROMPT string = "You are a helpful assistant who writes pull request bodies. You will be given information related to the pull request and you should use it to create a pull request body. It should detail the changes made to complete the pull request. Do not include file names. Make sure it details the main changes made, ignore any minor changes."
 
 var COMPRESS_DIFF_PROMPT string = "You are a helpful assistant who describes git diff changes. You will be given a Git diff and you should use it to create a description of the changes. The description should be no longer than 75 characters long and should describe the changes in the diff. Do not include the file names in the description."
+
+var REFACTOR_CODE_PROMPT string = `You are a helpful assistant who refactors code. You will be given a file what you are trying to accomplish in the refactor and you should refactor it to the best of your abilities. You should only output the code and should not output any other text.\n`

pkg/textfile/insert.go

@@ -38,3 +38,21 @@ func InsertLinesAtIndices(file string, indices []int, linesToInsert []string) (s
 
 	return strings.Join(lines, "\n"), nil
 }
+
+func ReplaceLines(code string, startLine int, endLine int, newText string) (string, error) {
+	if startLine < 1 || endLine < 1 {
+		return "", fmt.Errorf("startLine and endLine must be greater than or equal to 1")
+	}
+
+	if startLine > endLine {
+		return "", fmt.Errorf("startLine must be less than or equal to endLine")
+	}
+
+	lines := strings.Split(code, "\n")
+	if endLine > len(lines) {
+		return "", fmt.Errorf("endLine is greater than the number of lines in the code")
+	}
+
+	lines = append(lines[:startLine-1], append([]string{newText}, lines[endLine:]...)...)
+	return strings.Join(lines, "\n"), nil
+}

pkg/textfile/insert_test.go

@@ -0,0 +1,117 @@
+package textfile
+
+import (
+	"testing"
+)
+
+func TestReplaceLines(t *testing.T) {
+	testCases := []struct {
+		name           string
+		code           string
+		startLine      int
+		endLine        int
+		newText        string
+		expectedResult string
+		expectError    bool
+	}{
+		{
+			name:           "Replace single line",
+			code:           "Line 1\nLine 2\nLine 3",
+			startLine:      2,
+			endLine:        2,
+			newText:        "New Line 2",
+			expectedResult: "Line 1\nNew Line 2\nLine 3",
+			expectError:    false,
+		},
+		{
+			name:           "Replace multiple lines",
+			code:           "Line 1\nLine 2\nLine 3",
+			startLine:      1,
+			endLine:        2,
+			newText:        "New Line 1 and 2",
+			expectedResult: "New Line 1 and 2\nLine 3",
+			expectError:    false,
+		},
+		{
+			name:           "Invalid startLine and endLine",
+			code:           "Line 1\nLine 2\nLine 3",
+			startLine:      3,
+			endLine:        1,
+			newText:        "Invalid Test",
+			expectedResult: "",
+			expectError:    true,
+		},
+	}
+
+	for _, tc := range testCases {
+		t.Run(tc.name, func(t *testing.T) {
+			result, err := ReplaceLines(tc.code, tc.startLine, tc.endLine, tc.newText)
+
+			if tc.expectError && err == nil {
+				t.Errorf("Expected error, but got nil")
+			}
+
+			if !tc.expectError && err != nil {
+				t.Errorf("Unexpected error: %v", err)
+			}
+
+			if result != tc.expectedResult {
+				t.Errorf("Expected result: %q, but got: %q", tc.expectedResult, result)
+			}
+		})
+	}
+}
+
+func TestInsertLine(t *testing.T) {
+	testCases := []struct {
+		name           string
+		code           string
+		lineNumber     int
+		newText        string
+		expectedResult string
+		expectError    bool
+	}{
+		{
+			name:           "Insert at beginning",
+			code:           "Line 1\nLine 2\nLine 3",
+			lineNumber:     1,
+			newText:        "New Line 0",
+			expectedResult: "New Line 0\nLine 1\nLine 2\nLine 3",
+			expectError:    false,
+		},
+		{
+			name:           "Insert in the middle",
+			code:           "Line 1\nLine 2\nLine 3",
+			lineNumber:     2,
+			newText:        "New Line 2",
+			expectedResult: "Line 1\nNew Line 2\nLine 2\nLine 3",
+			expectError:    false,
+		},
+		{
+			name:           "Invalid line number",
+			code:           "Line 1\nLine 2\nLine 3",
+			lineNumber:     0,
+			newText:        "Invalid Test",
+			expectedResult: "",
+			expectError:    true,
+		},
+	}
+
+	for _, tc := range testCases {
+		t.Run(tc.name, func(t *testing.T) {
+			result, err := InsertLine(tc.code, tc.lineNumber, tc.newText)
+
+			if tc.expectError && err == nil {
+				t.Errorf("Expected error, but got nil")
+			}
+
+			if !tc.expectError && err != nil {
+				t.Errorf("Unexpected error: %v", err)
+			}
+
+			if result != tc.expectedResult {
+				t.Errorf("Expected result: %q, but got: %q", tc.expectedResult, result)
+			}
+		})
+	}
+}