Remove repeated code, make API consistent across functions

Author: chand1012

ai/calc_tokens.go

@@ -0,0 +1,16 @@
+package ai
+
+import "github.com/pandodao/tokenizer-go"
+
+func CalcTokens(inputs ...string) (int, error) {
+	total := int(0)
+	for _, input := range inputs {
+		tokens, err := tokenizer.CalToken(input)
+		if err != nil {
+			return -1, err
+		}
+		total += tokens
+	}
+
+	return total, nil
+}

ai/markdown.go

@@ -1,34 +1,20 @@
 package ai
 
 import (
-	"errors"
 	"fmt"
-	"os"
 	"strings"
-	"unicode/utf8"
 
 	gopenai "github.com/CasualCodersProjects/gopenai"
 	ai_types "github.com/CasualCodersProjects/gopenai/types"
 	"github.com/chand1012/ottodocs/constants"
-	"github.com/pandodao/tokenizer-go"
 )
 
-func Markdown(filePath, chatPrompt, APIKey string) (string, error) {
+func Markdown(filePath, contents, chatPrompt, APIKey string) (string, error) {
 	openai := gopenai.NewOpenAI(&gopenai.OpenAIOpts{
 		APIKey: APIKey,
 	})
 
-	// load the file
-	contents, err := os.ReadFile(filePath)
-	if err != nil {
-		return "", err
-	}
-
-	if !utf8.Valid(contents) {
-		return "", errors.New("the file is not valid UTF-8")
-	}
-
-	question := chatPrompt + "\n\n" + strings.TrimRight(string(contents), " \n")
+	question := chatPrompt + "\n\n" + strings.TrimRight(contents, " \n")
 
 	messages := []ai_types.ChatMessage{
 		{
@@ -41,7 +27,10 @@ func Markdown(filePath, chatPrompt, APIKey string) (string, error) {
 		},
 	}
 
-	tokens := tokenizer.MustCalToken(messages[0].Content) + tokenizer.MustCalToken(messages[1].Content)
+	tokens, err := CalcTokens(messages[0].Content, messages[1].Content)
+	if err != nil {
+		return "", fmt.Errorf("could not calculate tokens: %s", err)
+	}
 
 	maxTokens := constants.OPENAI_MAX_TOKENS - tokens
 

ai/question.go

@@ -7,7 +7,6 @@ import (
 	gopenai "github.com/CasualCodersProjects/gopenai"
 	ai_types "github.com/CasualCodersProjects/gopenai/types"
 	"github.com/chand1012/ottodocs/constants"
-	"github.com/pandodao/tokenizer-go"
 )
 
 func Question(filePath, fileContent, chatPrompt, APIKey string) (string, error) {
@@ -28,7 +27,10 @@ func Question(filePath, fileContent, chatPrompt, APIKey string) (string, error)
 		},
 	}
 
-	tokens := tokenizer.MustCalToken(messages[0].Content) + tokenizer.MustCalToken(messages[1].Content)
+	tokens, err := CalcTokens(messages[0].Content, messages[1].Content)
+	if err != nil {
+		return "", fmt.Errorf("could not calculate tokens: %s", err)
+	}
 
 	maxTokens := constants.OPENAI_MAX_TOKENS - tokens
 

ai/single_file.go

@@ -1,17 +1,13 @@
 package ai
 
 import (
-	"errors"
 	"fmt"
-	"os"
 	"path/filepath"
 	"strconv"
 	"strings"
-	"unicode/utf8"
 
 	gopenai "github.com/CasualCodersProjects/gopenai"
 	ai_types "github.com/CasualCodersProjects/gopenai/types"
-	"github.com/pandodao/tokenizer-go"
 
 	"github.com/chand1012/ottodocs/constants"
 	"github.com/chand1012/ottodocs/utils/textfile"
@@ -38,30 +34,20 @@ 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, chatPrompt, APIKey string) (string, error) {
+func SingleFile(filePath, contents, chatPrompt, APIKey string) (string, error) {
 
 	openai := gopenai.NewOpenAI(&gopenai.OpenAIOpts{
 		APIKey: APIKey,
 	})
 
-	// load the file
-	contents, err := os.ReadFile(filePath)
-	if err != nil {
-		return "", err
-	}
-
-	if !utf8.Valid(contents) {
-		return "", errors.New("the file is not valid UTF-8")
-	}
-
 	fileEnding := filepath.Ext(filePath)
 
 	commentOperator, ok := constants.CommentOperators[fileEnding]
 	if !ok {
 		return "", fmt.Errorf("the file type %s is not supported", fileEnding)
 	}
 
-	question := chatPrompt + "\n\n" + strings.TrimRight(string(contents), " \n")
+	question := chatPrompt + "\n\n" + strings.TrimRight(contents, " \n")
 
 	messages := []ai_types.ChatMessage{
 		{
@@ -74,7 +60,10 @@ func SingleFile(filePath, chatPrompt, APIKey string) (string, error) {
 		},
 	}
 
-	tokens := tokenizer.MustCalToken(messages[0].Content) + tokenizer.MustCalToken(messages[1].Content)
+	tokens, err := CalcTokens(messages[0].Content, messages[1].Content)
+	if err != nil {
+		return "", fmt.Errorf("could not calculate tokens: %s", err)
+	}
 
 	maxTokens := constants.OPENAI_MAX_TOKENS - tokens
 

cmd/ask.go

@@ -1,5 +1,5 @@
 /*
-Copyright © 2023 NAME HERE <EMAIL ADDRESS>
+Copyright © 2023 Chandler <chandler@chand1012.dev>
 */
 package cmd
 
@@ -10,7 +10,6 @@ import (
 	"strings"
 
 	"github.com/blevesearch/bleve/v2"
-	"github.com/chand1012/git2gpt/prompt"
 	"github.com/spf13/cobra"
 
 	"github.com/chand1012/ottodocs/ai"
@@ -22,7 +21,8 @@ import (
 var askCmd = &cobra.Command{
 	Use:   "ask",
 	Short: "Ask a question about a file or repo",
-	Long:  `Uses a vector search to find the most similar questions and answers`,
+	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.`,
 	Args: cobra.PositionalArgs(func(cmd *cobra.Command, args []string) error {
 		if len(args) < 1 {
 			return fmt.Errorf("requires a path to a repository or file")
@@ -47,15 +47,6 @@ var askCmd = &cobra.Command{
 			chatPrompt = strings.TrimRight(chatPrompt, " \n")
 		}
 
-		// if .index.bleve exists, delete it
-		if _, err := os.Stat(filepath.Join(args[0], ".index.bleve")); err == nil {
-			err = os.RemoveAll(filepath.Join(args[0], ".index.bleve"))
-			if err != nil {
-				log.Errorf("Error deleting index: %s", err)
-				os.Exit(1)
-			}
-		}
-
 		info, err := os.Stat(repoPath)
 		if err != nil {
 			log.Errorf("Error getting file info: %s", err)
@@ -78,13 +69,12 @@ var askCmd = &cobra.Command{
 				os.Exit(1)
 			}
 			// index the repo
-			ignoreList := prompt.GenerateIgnoreList(repoPath, ignoreFilePath, !ignoreGitignore)
-			ignoreList = append(ignoreList, filepath.Join(repoPath, ".gptignore"))
-			repo, err := prompt.ProcessGitRepo(repoPath, ignoreList)
+			repo, err := utils.GetRepo(repoPath, ignoreFilePath, ignoreGitignore)
 			if err != nil {
 				log.Errorf("Error processing repo: %s", err)
 				os.Exit(1)
 			}
+
 			// index the files
 			for _, file := range repo.Files {
 				err = index.Index(file.Path, file)
@@ -139,21 +129,29 @@ var askCmd = &cobra.Command{
 
 		fmt.Println("Asking question about " + fileName + "...")
 
-		// load the file
-		content, err := os.ReadFile(fileName)
+		content, err := utils.LoadFile(fileName)
 		if err != nil {
-			log.Errorf("Error reading file: %s", err)
+			log.Errorf("Error loading file: %s", err)
 			os.Exit(1)
 		}
 
-		resp, err := ai.Question(fileName, string(content), chatPrompt, conf.APIKey)
+		resp, err := ai.Question(fileName, content, chatPrompt, conf.APIKey)
 
 		if err != nil {
 			log.Errorf("Error asking question: %s", err)
 			os.Exit(1)
 		}
 
 		fmt.Println(resp)
+
+		// if .index.bleve exists, delete it
+		if _, err := os.Stat(filepath.Join(args[0], ".index.bleve")); err == nil {
+			err = os.RemoveAll(filepath.Join(args[0], ".index.bleve"))
+			if err != nil {
+				log.Errorf("Error deleting index: %s", err)
+				os.Exit(1)
+			}
+		}
 	},
 }
 

cmd/doc.go

@@ -11,6 +11,7 @@ import (
 
 	"github.com/chand1012/ottodocs/ai"
 	"github.com/chand1012/ottodocs/config"
+	"github.com/chand1012/ottodocs/utils"
 )
 
 // docCmd represents the doc command
@@ -39,10 +40,16 @@ var docCmd = &cobra.Command{
 
 		var contents string
 
+		fileContents, err := utils.LoadFile(filePath)
+		if err != nil {
+			log.Errorf("Error: %s", err)
+			os.Exit(1)
+		}
+
 		if inlineMode || !markdownMode {
-			contents, err = ai.SingleFile(filePath, chatPrompt, conf.APIKey)
+			contents, err = ai.SingleFile(filePath, fileContents, chatPrompt, conf.APIKey)
 		} else {
-			contents, err = ai.Markdown(filePath, chatPrompt, conf.APIKey)
+			contents, err = ai.Markdown(filePath, fileContents, chatPrompt, conf.APIKey)
 		}
 
 		if err != nil {

cmd/docs.go

@@ -8,9 +8,9 @@ import (
 	"os"
 	"path/filepath"
 
-	"github.com/chand1012/git2gpt/prompt"
 	"github.com/chand1012/ottodocs/ai"
 	"github.com/chand1012/ottodocs/config"
+	"github.com/chand1012/ottodocs/utils"
 	"github.com/spf13/cobra"
 )
 
@@ -61,9 +61,7 @@ search for files in the directory and document them.
 			os.Exit(1)
 		}
 
-		ignoreList := prompt.GenerateIgnoreList(repoPath, ignoreFilePath, !ignoreGitignore)
-		ignoreList = append(ignoreList, filepath.Join(repoPath, ".gptignore"))
-		repo, err := prompt.ProcessGitRepo(repoPath, ignoreList)
+		repo, err := utils.GetRepo(repoPath, ignoreFilePath, ignoreGitignore)
 		if err != nil {
 			log.Errorf("Error: %s", err)
 			os.Exit(1)
@@ -82,10 +80,16 @@ search for files in the directory and document them.
 				chatPrompt = "Write documentation for the following code snippet. The file name is" + file.Path + ":"
 			}
 
+			fileContents, err := utils.LoadFile(path)
+			if err != nil {
+				log.Warnf("Error loading file %s: %s", path, err)
+				continue
+			}
+
 			if inlineMode || !markdownMode {
-				contents, err = ai.SingleFile(path, chatPrompt, conf.APIKey)
+				contents, err = ai.SingleFile(path, fileContents, chatPrompt, conf.APIKey)
 			} else {
-				contents, err = ai.Markdown(path, chatPrompt, conf.APIKey)
+				contents, err = ai.Markdown(path, fileContents, chatPrompt, conf.APIKey)
 			}
 
 			if err != nil {

docs/otto_ask.md

@@ -0,0 +1,25 @@
+## otto ask
+
+Ask a question about a file or repo
+
+### Synopsis
+
+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.
+
+```
+otto ask [flags]
+```
+
+### Options
+
+```
+  -h, --help               help for ask
+  -n, --ignore string      path to .gptignore file
+  -g, --ignore-gitignore   ignore .gitignore file
+  -q, --question string    The question to ask
+```
+
+### SEE ALSO
+
+* [otto](otto.md)	 - Document your code with ease

utils/get_repo.go

@@ -0,0 +1,18 @@
+package utils
+
+import (
+	"path/filepath"
+
+	"github.com/chand1012/git2gpt/prompt"
+)
+
+// GetRepo returns a GitRepo object for the given repoPath
+func GetRepo(repoPath, ignoreFilePath string, ignoreGitIgnore bool) (*prompt.GitRepo, error) {
+	ignoreList := prompt.GenerateIgnoreList(repoPath, ignoreFilePath, !ignoreGitIgnore)
+	ignoreList = append(ignoreList, filepath.Join(repoPath, ".gptignore"))
+	repo, err := prompt.ProcessGitRepo(repoPath, ignoreList)
+	if err != nil {
+		return nil, err
+	}
+	return repo, nil
+}

utils/load_file.go

@@ -0,0 +1,22 @@
+package utils
+
+import (
+	"errors"
+	"os"
+	"unicode/utf8"
+)
+
+// Loads a file and return an error if its not valid UTF-8
+func LoadFile(filePath string) (string, error) {
+	// load the file
+	contents, err := os.ReadFile(filePath)
+	if err != nil {
+		return "", err
+	}
+
+	if !utf8.Valid(contents) {
+		return "", errors.New("the file is not valid UTF-8")
+	}
+
+	return string(contents), nil
+}