Refactor AI functions, add commit command

chand1012 · chand1012 · commit faf3f2814a24 · 2023-04-21T16:19:17.000-04:00
diff --git a/cmd/ask.go b/cmd/ask.go
@@ -105,7 +105,7 @@ Requires a path to a repository or file as a positional argument.`,
 			os.Exit(1)
 		}
 
-		resp, err := ai.Question(fileName, content, chatPrompt, conf.APIKey, conf.Model)
+		resp, err := ai.Question(fileName, content, chatPrompt, conf)
 
 		if err != nil {
 			log.Errorf("Error asking question: %s", err)
diff --git a/cmd/cmd.go b/cmd/cmd.go
@@ -44,7 +44,7 @@ This command is only supported on MacOS and Linux using Bash or Zsh. Windows and
 
 		// fmt.Println("History:", history)
 
-		resp, err := ai.CmdQuestion(history, chatPrompt, conf.APIKey, conf.Model)
+		resp, err := ai.CmdQuestion(history, chatPrompt, conf)
 
 		if err != nil {
 			log.Error(err)
diff --git a/cmd/commit.go b/cmd/commit.go
@@ -0,0 +1,61 @@
+/*
+Copyright © 2023 Chandler <chandler@chand1012.dev>
+*/
+package cmd
+
+import (
+	"fmt"
+	"os"
+	"os/exec"
+
+	"github.com/chand1012/ottodocs/pkg/ai"
+	"github.com/chand1012/ottodocs/pkg/config"
+	"github.com/spf13/cobra"
+)
+
+// 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.`,
+	Run: func(cmd *cobra.Command, args []string) {
+		conf, err := config.Load()
+		if err != nil || conf.APIKey == "" {
+			// if the API key is not set, prompt the user to login
+			log.Error("Please login first.")
+			log.Error("Run `ottodocs login` to login.")
+			os.Exit(1)
+		}
+
+		c := exec.Command("git", "diff")
+		diffBytes, err := c.Output()
+		if err != nil {
+			log.Error(err)
+			os.Exit(1)
+		}
+
+		diff := string(diffBytes)
+
+		msg, err := ai.CommitMessage(diff, conventional, conf)
+		if err != nil {
+			log.Error(err)
+			os.Exit(1)
+		}
+
+		fmt.Println("Commit message:", msg)
+	},
+}
+
+func init() {
+	RootCmd.AddCommand(commitCmd)
+
+	// Here you will define your flags and configuration settings.
+
+	// Cobra supports Persistent Flags which will work for this command
+	// and all subcommands, e.g.:
+	// commitCmd.PersistentFlags().String("foo", "", "A help for foo")
+
+	// Cobra supports local flags which will only run when this command
+	// is called directly, e.g.:
+	commitCmd.Flags().BoolVarP(&conventional, "conventional", "c", false, "use conventional commits")
+}
diff --git a/cmd/docs.go b/cmd/docs.go
@@ -89,9 +89,9 @@ search for files in the directory and document them. If a single file is specifi
 				}
 
 				if inlineMode || !markdownMode {
-					contents, err = ai.SingleFile(path, fileContents, chatPrompt, conf.APIKey, conf.Model)
+					contents, err = ai.SingleFile(path, fileContents, chatPrompt, conf)
 				} else {
-					contents, err = ai.Markdown(path, fileContents, chatPrompt, conf.APIKey, conf.Model)
+					contents, err = ai.Markdown(path, fileContents, chatPrompt, conf)
 				}
 
 				if err != nil {
@@ -152,9 +152,9 @@ search for files in the directory and document them. If a single file is specifi
 			}
 
 			if inlineMode || !markdownMode {
-				contents, err = ai.SingleFile(filePath, fileContents, chatPrompt, conf.APIKey, conf.Model)
+				contents, err = ai.SingleFile(filePath, fileContents, chatPrompt, conf)
 			} else {
-				contents, err = ai.Markdown(filePath, fileContents, chatPrompt, conf.APIKey, conf.Model)
+				contents, err = ai.Markdown(filePath, fileContents, chatPrompt, conf)
 			}
 
 			if err != nil {
diff --git a/cmd/vars.go b/cmd/vars.go
@@ -20,6 +20,8 @@ var inlineMode bool
 var markdownMode bool
 var overwriteOriginal bool
 
+var conventional bool // use conventional commits
+
 var log = l.NewWithOptions(os.Stderr, l.Options{
 	Level:           l.InfoLevel,
 	ReportTimestamp: false,
diff --git a/pkg/ai/cmd_question.go b/pkg/ai/cmd_question.go
@@ -6,12 +6,13 @@ import (
 	gopenai "github.com/CasualCodersProjects/gopenai"
 	ai_types "github.com/CasualCodersProjects/gopenai/types"
 	"github.com/chand1012/ottodocs/pkg/calc"
+	"github.com/chand1012/ottodocs/pkg/config"
 	"github.com/chand1012/ottodocs/pkg/constants"
 )
 
-func CmdQuestion(history []string, chatPrompt, APIKey, model string) (string, error) {
+func CmdQuestion(history []string, chatPrompt string, conf *config.Config) (string, error) {
 	openai := gopenai.NewOpenAI(&gopenai.OpenAIOpts{
-		APIKey: APIKey,
+		APIKey: conf.APIKey,
 	})
 
 	questionNoHistory := "\nQuestion: " + chatPrompt + "\n\nAnswer:"
@@ -24,7 +25,7 @@ func CmdQuestion(history []string, chatPrompt, APIKey, model string) (string, er
 	for i := len(history) - 1; i >= 0; i-- {
 		newHistory := history[i] + "\n"
 		tokens := calc.EstimateTokens(newHistory) + qTokens + calc.EstimateTokens(historyQuestion) + commandPromptTokens
-		if tokens < constants.OPENAI_MAX_TOKENS {
+		if tokens < calc.GetMaxTokens(conf.Model) {
 			historyQuestion += newHistory
 		} else {
 			break
@@ -48,16 +49,16 @@ func CmdQuestion(history []string, chatPrompt, APIKey, model string) (string, er
 
 	tokens := calc.EstimateTokens(messages[0].Content, messages[1].Content)
 
-	maxTokens := constants.OPENAI_MAX_TOKENS - tokens
+	maxTokens := calc.GetMaxTokens(conf.Model) - tokens
 
 	if maxTokens < 0 {
-		return "", fmt.Errorf("the prompt is too long. max length is %d. Got %d", constants.OPENAI_MAX_TOKENS, tokens)
+		return "", fmt.Errorf("the prompt is too long. max length is %d. Got %d", calc.GetMaxTokens(conf.Model), tokens)
 	}
 
 	req := ai_types.NewDefaultChatRequest("")
 	req.Messages = messages
 	req.MaxTokens = maxTokens
-	req.Model = model
+	req.Model = conf.Model
 	// lower the temperature to make the model more deterministic
 	// req.Temperature = 0.3
 
diff --git a/pkg/ai/commit.go b/pkg/ai/commit.go
@@ -0,0 +1,58 @@
+package ai
+
+import (
+	"fmt"
+
+	gopenai "github.com/CasualCodersProjects/gopenai"
+	ai_types "github.com/CasualCodersProjects/gopenai/types"
+	"github.com/chand1012/ottodocs/pkg/calc"
+	"github.com/chand1012/ottodocs/pkg/config"
+	"github.com/chand1012/ottodocs/pkg/constants"
+)
+
+func CommitMessage(diff string, conventional bool, conf *config.Config) (string, error) {
+	openai := gopenai.NewOpenAI(&gopenai.OpenAIOpts{
+		APIKey: conf.APIKey,
+	})
+
+	sysMessage := constants.GIT_DIFF_PROMPT_STD
+	if conventional {
+		sysMessage = constants.GIT_DIFF_PROMPT_CONVENTIONAL
+	}
+
+	messages := []ai_types.ChatMessage{
+		{
+			Content: sysMessage,
+			Role:    "system",
+		},
+		{
+			Content: diff,
+			Role:    "user",
+		},
+	}
+
+	tokens, err := calc.PreciseTokens(messages[0].Content, messages[1].Content)
+	if err != nil {
+		return "", fmt.Errorf("could not calculate tokens: %s", err)
+	}
+
+	maxTokens := calc.GetMaxTokens(conf.Model) - tokens
+
+	if maxTokens < 0 {
+		return "", fmt.Errorf("the prompt is too long. max length is %d. Got %d", calc.GetMaxTokens(conf.Model), tokens)
+	}
+
+	req := ai_types.NewDefaultChatRequest("")
+	req.Messages = messages
+	req.MaxTokens = maxTokens
+	req.Model = conf.Model
+
+	resp, err := openai.CreateChat(req)
+	if err != nil {
+		return "", err
+	}
+
+	message := resp.Choices[0].Message.Content
+
+	return message, nil
+}
diff --git a/pkg/ai/markdown.go b/pkg/ai/markdown.go
@@ -7,12 +7,13 @@ import (
 	gopenai "github.com/CasualCodersProjects/gopenai"
 	ai_types "github.com/CasualCodersProjects/gopenai/types"
 	"github.com/chand1012/ottodocs/pkg/calc"
+	"github.com/chand1012/ottodocs/pkg/config"
 	"github.com/chand1012/ottodocs/pkg/constants"
 )
 
-func Markdown(filePath, contents, chatPrompt, APIKey, model string) (string, error) {
+func Markdown(filePath, contents, chatPrompt string, conf *config.Config) (string, error) {
 	openai := gopenai.NewOpenAI(&gopenai.OpenAIOpts{
-		APIKey: APIKey,
+		APIKey: conf.APIKey,
 	})
 
 	question := chatPrompt + "\n\n" + strings.TrimRight(contents, " \n")
@@ -33,16 +34,16 @@ func Markdown(filePath, contents, chatPrompt, APIKey, model string) (string, err
 		return "", fmt.Errorf("could not calculate tokens: %s", err)
 	}
 
-	maxTokens := constants.OPENAI_MAX_TOKENS - tokens
+	maxTokens := calc.GetMaxTokens(conf.Model) - tokens
 
 	if maxTokens < 0 {
-		return "", fmt.Errorf("the prompt is too long. max length is %d. Got %d", constants.OPENAI_MAX_TOKENS, tokens)
+		return "", fmt.Errorf("the prompt is too long. max length is %d. Got %d", calc.GetMaxTokens(conf.Model), tokens)
 	}
 
 	req := ai_types.NewDefaultChatRequest("")
 	req.Messages = messages
 	req.MaxTokens = maxTokens
-	req.Model = model
+	req.Model = conf.Model
 	// lower the temperature to make the model more deterministic
 	// req.Temperature = 0.3
 
diff --git a/pkg/ai/question.go b/pkg/ai/question.go
@@ -7,12 +7,13 @@ import (
 	gopenai "github.com/CasualCodersProjects/gopenai"
 	ai_types "github.com/CasualCodersProjects/gopenai/types"
 	"github.com/chand1012/ottodocs/pkg/calc"
+	"github.com/chand1012/ottodocs/pkg/config"
 	"github.com/chand1012/ottodocs/pkg/constants"
 )
 
-func Question(filePath, fileContent, chatPrompt, APIKey, model string) (string, error) {
+func Question(filePath, fileContent, chatPrompt string, conf *config.Config) (string, error) {
 	openai := gopenai.NewOpenAI(&gopenai.OpenAIOpts{
-		APIKey: APIKey,
+		APIKey: conf.APIKey,
 	})
 
 	question := "File Name: " + filePath + "\nQuestion: " + chatPrompt + "\n\n" + strings.TrimRight(string(fileContent), " \n") + "\nAnswer:"
@@ -33,16 +34,16 @@ func Question(filePath, fileContent, chatPrompt, APIKey, model string) (string,
 		return "", fmt.Errorf("could not calculate tokens: %s", err)
 	}
 
-	maxTokens := constants.OPENAI_MAX_TOKENS - tokens
+	maxTokens := calc.GetMaxTokens(conf.Model) - tokens
 
 	if maxTokens < 0 {
-		return "", fmt.Errorf("the prompt is too long. max length is %d. Got %d", constants.OPENAI_MAX_TOKENS, tokens)
+		return "", fmt.Errorf("the prompt is too long. max length is %d. Got %d", calc.GetMaxTokens(conf.Model), tokens)
 	}
 
 	req := ai_types.NewDefaultChatRequest("")
 	req.Messages = messages
 	req.MaxTokens = maxTokens
-	req.Model = model
+	req.Model = conf.Model
 	// lower the temperature to make the model more deterministic
 	// req.Temperature = 0.3
 
diff --git a/pkg/ai/single_file.go b/pkg/ai/single_file.go
@@ -10,6 +10,7 @@ import (
 	ai_types "github.com/CasualCodersProjects/gopenai/types"
 
 	"github.com/chand1012/ottodocs/pkg/calc"
+	"github.com/chand1012/ottodocs/pkg/config"
 	"github.com/chand1012/ottodocs/pkg/constants"
 	"github.com/chand1012/ottodocs/pkg/textfile"
 )
@@ -35,10 +36,10 @@ func extractLineNumber(line string) (int, error) {
 }
 
 // Document a file using the OpenAI ChatGPT API. Takes a file path, a prompt, and an API key as arguments.
-func SingleFile(filePath, contents, chatPrompt, APIKey, model string) (string, error) {
+func SingleFile(filePath, contents, chatPrompt string, conf *config.Config) (string, error) {
 
 	openai := gopenai.NewOpenAI(&gopenai.OpenAIOpts{
-		APIKey: APIKey,
+		APIKey: conf.APIKey,
 	})
 
 	fileEnding := filepath.Ext(filePath)
@@ -66,16 +67,16 @@ func SingleFile(filePath, contents, chatPrompt, APIKey, model string) (string, e
 		return "", fmt.Errorf("could not calculate tokens: %s", err)
 	}
 
-	maxTokens := constants.OPENAI_MAX_TOKENS - tokens
+	maxTokens := calc.GetMaxTokens(conf.Model) - tokens
 
 	if maxTokens < 0 {
-		return "", fmt.Errorf("the prompt is too long. max length is %d. Got %d", constants.OPENAI_MAX_TOKENS, tokens)
+		return "", fmt.Errorf("the prompt is too long. max length is %d. Got %d", calc.GetMaxTokens(conf.Model), tokens)
 	}
 
 	req := ai_types.NewDefaultChatRequest("")
 	req.Messages = messages
 	req.MaxTokens = maxTokens
-	req.Model = model
+	req.Model = conf.Model
 	// lower the temperature to make the model more deterministic
 	req.Temperature = 0.3
 
diff --git a/pkg/calc/tokens.go b/pkg/calc/tokens.go
@@ -1,6 +1,10 @@
 package calc
 
-import "github.com/pandodao/tokenizer-go"
+import (
+	"strings"
+
+	"github.com/pandodao/tokenizer-go"
+)
 
 // Precise token count
 func PreciseTokens(inputs ...string) (int, error) {
@@ -26,3 +30,14 @@ func EstimateTokens(inputs ...string) int {
 
 	return total
 }
+
+func GetMaxTokens(model string) int {
+	if strings.Contains(model, "32k") {
+		return 32000
+	}
+	if strings.Contains(model, "4") {
+		return 8000
+	}
+
+	return 4000
+}
diff --git a/pkg/constants/openAI.go b/pkg/constants/openAI.go
diff --git a/pkg/constants/prompts.go b/pkg/constants/prompts.go
@@ -29,3 +29,7 @@ var COMMAND_QUESTION_PROMPT string = `You are a helpful assistant who answers qu
 - The answer must be in English.
 - If there is no way to answer the question, you should say so.
 - The answer must be AT LEAST one sentence long.`
+
+var GIT_DIFF_PROMPT_STD string = `You are a helpful assistant who writes git commit messages. You will be given a Git diff and you should use it to create a commit message. The commit message should be roughly 50 characters long and should describe the changes in the diff. Do not include the file names in the commit message.`
+
+var GIT_DIFF_PROMPT_CONVENTIONAL string = `You are a helpful assistant who writes git commit messages. You will be given a Git diff and you should use it to create a commit message. The commit message should be roughly 50 characters long and should describe the changes in the diff. Do not include the file names in the commit message. The commit message should follow the conventional commit format.`
diff --git a/pkg/textfile/split.go b/pkg/textfile/split.go
@@ -5,7 +5,6 @@ import (
 	"strings"
 
 	"github.com/chand1012/ottodocs/pkg/calc"
-	"github.com/chand1012/ottodocs/pkg/constants"
 	"github.com/chand1012/ottodocs/pkg/utils"
 )
 
@@ -21,11 +20,8 @@ func (s *SplitFile) Hash() string {
 	return fmt.Sprintf("%s#%d-%d", s.Path, s.StartLine, s.EndLine)
 }
 
-func Split(path string) ([]SplitFile, error) {
-	// read the file with the os package
-	// split the file into lines with a max of 4000 tokens
-	// return a slice of SplitFile structs
-
+// splits the file based on the max tokens allowed by the model
+func Split(path, model string) ([]SplitFile, error) {
 	contents, err := utils.LoadFile(path)
 	if err != nil {
 		return nil, fmt.Errorf("could not load file: %s", err)
@@ -48,7 +44,7 @@ func Split(path string) ([]SplitFile, error) {
 
 		lastSplitFile := splitFiles[len(splitFiles)-1]
 
-		if calc.EstimateTokens(lastSplitFile.Contents)+calc.EstimateTokens(line) > constants.OPENAI_MAX_TOKENS {
+		if calc.EstimateTokens(lastSplitFile.Contents)+calc.EstimateTokens(line) > calc.GetMaxTokens(model) {
 			splitFiles = append(splitFiles, SplitFile{
 				Path:      path,
 				Contents:  line,

Original file line number	Diff line number	Diff line change
@@ -105,7 +105,7 @@ Requires a path to a repository or file as a positional argument.`,
`105`	`105`	`os.Exit(1)`
`106`	`106`	`}`
`107`	`107`
`108`		`- resp, err := ai.Question(fileName, content, chatPrompt, conf.APIKey, conf.Model)`
	`108`	`+ resp, err := ai.Question(fileName, content, chatPrompt, conf)`
`109`	`109`
`110`	`110`	`if err != nil {`
`111`	`111`	`log.Errorf("Error asking question: %s", err)`
Original file line number	Diff line number	Diff line change
`@@ -89,9 +89,9 @@ search for files in the directory and document them. If a single file is specifi`
`89`	`89`	`}`
`90`	`90`
`91`	`91`	`if inlineMode \|\| !markdownMode {`
`92`		`- contents, err = ai.SingleFile(path, fileContents, chatPrompt, conf.APIKey, conf.Model)`
	`92`	`+ contents, err = ai.SingleFile(path, fileContents, chatPrompt, conf)`
`93`	`93`	`} else {`
`94`		`- contents, err = ai.Markdown(path, fileContents, chatPrompt, conf.APIKey, conf.Model)`
	`94`	`+ contents, err = ai.Markdown(path, fileContents, chatPrompt, conf)`
`95`	`95`	`}`
`96`	`96`
`97`	`97`	`if err != nil {`
`@@ -152,9 +152,9 @@ search for files in the directory and document them. If a single file is specifi`
`152`	`152`	`}`
`153`	`153`
`154`	`154`	`if inlineMode \|\| !markdownMode {`
`155`		`- contents, err = ai.SingleFile(filePath, fileContents, chatPrompt, conf.APIKey, conf.Model)`
	`155`	`+ contents, err = ai.SingleFile(filePath, fileContents, chatPrompt, conf)`
`156`	`156`	`} else {`
`157`		`- contents, err = ai.Markdown(filePath, fileContents, chatPrompt, conf.APIKey, conf.Model)`
	`157`	`+ contents, err = ai.Markdown(filePath, fileContents, chatPrompt, conf)`
`158`	`158`	`}`
`159`	`159`
`160`	`160`	`if err != nil {`