Add context files support to edit and issue commands

Author: chand1012

cmd/edit.go

@@ -86,9 +86,21 @@ Example: otto edit main.go --start 1 --end 10 --goal "Refactor the function"`,
 
 		var prompt string
 		if editCode != "" {
-			prompt = constants.EDIT_CODE_PROMPT + "\nEDIT: " + editCode + "\n\nGOAL: " + chatPrompt + "\n\nFILE: " + contents
+			prompt = constants.EDIT_CODE_PROMPT + "\nEDIT: " + editCode + "\n\nGOAL: " + chatPrompt + "\n\nFILE: " + filePath + "\n\n" + contents
 		} else {
-			prompt = constants.EDIT_CODE_PROMPT + "\nGOAL: " + chatPrompt + "\n\nFILE: " + contents
+			prompt = constants.EDIT_CODE_PROMPT + "\nGOAL: " + chatPrompt + "\n\nFILE: " + filePath + "\n\n" + contents
+		}
+
+		if len(contextFiles) > 0 {
+			var contextContent string
+			for _, contextFile := range contextFiles {
+				contextContent, err = utils.LoadFile(contextFile)
+				if err != nil {
+					log.Errorf("Error loading context file: %s", err)
+					continue
+				}
+				prompt += "\n\nCONTEXT: " + contextFile + "\n\n" + contextContent
+			}
 		}
 
 		stream, err := ai.SimpleStreamRequest(prompt, c)
@@ -154,4 +166,5 @@ func init() {
 	editCmd.Flags().IntVarP(&startLine, "start", "s", 1, "Start line")
 	editCmd.Flags().IntVarP(&endLine, "end", "e", 0, "End line")
 	editCmd.Flags().StringVarP(&chatPrompt, "goal", "g", "", "Goal of the edit")
+	editCmd.Flags().StringSliceVarP(&contextFiles, "context", "c", []string{}, "Context files")
 }

cmd/issue.go

@@ -133,6 +133,26 @@ var issueCmd = &cobra.Command{
 				prompt += comment
 				tokens += commentTokens
 			}
+		} else if len(contextFiles) > 0 {
+			log.Debug("Using specified files...")
+			var contextContent string
+			for _, contextFile := range contextFiles {
+				contextContent, err = utils.LoadFile(contextFile)
+				if err != nil {
+					log.Errorf("Error loading context file: %s", err)
+					continue
+				}
+				prompt += "\n\nFile: " + contextFile + "\n\n" + contextContent
+				fileTokens, err := calc.PreciseTokens(contextContent)
+				if err != nil {
+					log.Errorf("Error getting tokens: %s", err)
+					os.Exit(1)
+				}
+				if tokens+fileTokens > calc.GetMaxTokens(c.Model) {
+					break
+				}
+				tokens += fileTokens
+			}
 		} else {
 			log.Debug("Using repo contents...")
 			// get the repo context here
@@ -246,6 +266,7 @@ func init() {
 
 	issueCmd.Flags().IntVarP(&issuePRNumber, "number", "n", 0, "the number of the issue to get")
 	issueCmd.Flags().StringVarP(&question, "question", "q", "", "the question to ask Otto")
+	issueCmd.Flags().StringSliceVarP(&contextFiles, "context", "f", []string{}, "the files to use as context")
 	issueCmd.Flags().BoolVarP(&useComments, "comments", "c", false, "use comments instead of git repo for context")
 	issueCmd.Flags().BoolVarP(&promptOnly, "prompt-only", "p", false, "only generate a prompt, don't ask Otto")
 	issueCmd.Flags().BoolVarP(&verbose, "verbose", "v", false, "verbose output")

cmd/vars.go

@@ -47,6 +47,8 @@ var appendFile bool
 var previousTag string
 var currentTag string
 
+var contextFiles []string
+
 var log = l.NewWithOptions(os.Stderr, l.Options{
 	Level:           l.InfoLevel,
 	ReportTimestamp: false,

pkg/constants/prompts.go

@@ -40,7 +40,7 @@ var PR_BODY_PROMPT string = "You are a helpful assistant who writes pull request
 
 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 EDIT_CODE_PROMPT string = `You are a helpful assistant who edits code. You will be given a file what you are trying to accomplish in the edit and you should edit it to the best of your abilities. Only edit the code you are told. All code to edit will be preceded by "EDIT:". If "EDIT:" is omitted, assume you must edit the entire file. The goal of the edit will be preceded by "GOAL:". The entire file will be preceded by "FILE:". If there is not content after "FILE:", assume you are writing new code. Make sure to use the language specified in the task. The output code should be unformatted. Use no markdown.\n`
+var EDIT_CODE_PROMPT string = `You are a helpful assistant who edits code. You will be given a file what you are trying to accomplish in the edit and you should edit it to the best of your abilities. Only edit the code you are told. All code to edit will be preceded by "EDIT:". If "EDIT:" is omitted, assume you must edit the entire file. The goal of the edit will be preceded by "GOAL:". The entire file will be preceded by "FILE:". If there is not content after "FILE:", assume you are writing new code. There may also be additional files added to give you more information about the project, those will be preceded by 'CONTEXT: '. DO NOT EDIT THOSE FILES. Make sure to use the language specified in the task. The output code should be unformatted. Use no markdown.\n`
 
 var RELEASE_PROMPT string = `You are a helpful assistant who creates GitHub release notes from git commit logs. You will be given a series of git commit messages and your task is to summarize these messages into release notes. The release notes should:
 - Be concise and informative about the changes made in the release.