feat: add support for customizable prompt templates from directories (#224)

- Add `promptFolder` variable and setting it via command line flag
- Create and validate `promptFolder` during configuration initialization
- Implement loading of custom prompt templates from `promptFolder`
- Implement the `prompt` command to load default prompt data
- Add function to get raw data of the template
- Initialize `templates` as an empty map to avoid redundant checks
- Add support to load templates from a specified directory
- Add tests for loading templates from a directory

Signed-off-by: appleboy <[email protected]>

Author: appleboy

cmd/cmd.go

@@ -21,8 +21,9 @@ var rootCmd = &cobra.Command{
 
 // Used for flags.
 var (
-	cfgFile  string
-	replacer = strings.NewReplacer("-", "_", ".", "_")
+	cfgFile      string
+	promptFolder string
+	replacer     = strings.NewReplacer("-", "_", ".", "_")
 )
 
 const (
@@ -34,12 +35,14 @@ func init() {
 	cobra.OnInitialize(initConfig)
 
 	rootCmd.PersistentFlags().StringVar(&cfgFile, "config", "", "config file (default is $HOME/.config/codegpt/.codegpt.yaml)")
+	rootCmd.PersistentFlags().StringVar(&promptFolder, "prompt_folder", "", "prompt folder (default is $HOME/.config/codegpt/prompt)")
 	rootCmd.AddCommand(versionCmd)
 	rootCmd.AddCommand(configCmd)
 	rootCmd.AddCommand(commitCmd)
 	rootCmd.AddCommand(hookCmd)
 	rootCmd.AddCommand(reviewCmd)
 	rootCmd.AddCommand(CompletionCmd)
+	rootCmd.AddCommand(promptCmd)
 
 	// hide completion command
 	rootCmd.CompletionOptions.HiddenDefaultCmd = true
@@ -75,6 +78,30 @@ func initConfig() {
 		}
 	}
 
+	if promptFolder != "" {
+		viper.Set("prompt_folder", promptFolder)
+		if file.IsFile(promptFolder) {
+			log.Fatalf("prompt folder %s is a file", promptFolder)
+		}
+		// create the prompt folder if it doesn't exist
+		if !file.IsDir(promptFolder) {
+			if err := os.MkdirAll(promptFolder, os.ModePerm); err != nil {
+				log.Fatal(err)
+			}
+		}
+	} else {
+		// Find home directory.
+		home, err := os.UserHomeDir()
+		cobra.CheckErr(err)
+		targetFolder := path.Join(home, ".config", "codegpt", "prompt")
+		if !file.IsDir(targetFolder) {
+			if err := os.MkdirAll(targetFolder, os.ModePerm); err != nil {
+				log.Fatal(err)
+			}
+		}
+		viper.Set("prompt_folder", targetFolder)
+	}
+
 	viper.AutomaticEnv()
 	viper.SetEnvKeyReplacer(replacer)
 

cmd/hepler.go

@@ -65,5 +65,13 @@ func check() error {
 		return fmt.Errorf("template variables file not found: %s", templateVarsFile)
 	}
 
+	// load custom prompt
+	promptFolder := viper.GetString("prompt_folder")
+	if promptFolder != "" {
+		if err := util.LoadTemplatesFromDir(promptFolder); err != nil {
+			return fmt.Errorf("failed to load custom prompt templates: %s", err)
+		}
+	}
+
 	return nil
 }

cmd/prompt.go

@@ -0,0 +1,71 @@
+package cmd
+
+import (
+	"os"
+	"path"
+
+	"github.com/appleboy/CodeGPT/prompt"
+	"github.com/erikgeiser/promptkit/confirmation"
+	"github.com/fatih/color"
+	"github.com/spf13/cobra"
+	"github.com/spf13/viper"
+)
+
+var loadPromptData bool
+
+func init() {
+	promptCmd.PersistentFlags().BoolVar(&loadPromptData, "load", false,
+		"load default prompt data")
+}
+
+var defaultPromptDataKeys = []string{
+	prompt.CodeReviewTemplate,
+	prompt.SummarizeFileDiffTemplate,
+	prompt.SummarizeTitleTemplate,
+	prompt.ConventionalCommitTemplate,
+}
+
+// promptCmd represents the command to load default prompt data.
+// It uses the "prompt" keyword and provides a short description: "load default prompt data".
+// The command executes the RunE function which checks if the loadPromptData flag is set.
+// If set, it prompts the user for confirmation to load the default prompt data, which will overwrite existing data.
+// Upon confirmation, it retrieves the prompt folder path from the configuration and saves the default prompt data keys to the specified folder.
+// If any error occurs during the process, it returns the error.
+var promptCmd = &cobra.Command{
+	Use:   "prompt",
+	Short: "load default prompt data",
+	RunE: func(cmd *cobra.Command, args []string) error {
+		if !loadPromptData {
+			return nil
+		}
+
+		confirm, err := confirmation.New("Do you want to load default prompt data, will overwrite your data", confirmation.No).RunPrompt()
+		if err != nil || !confirm {
+			return err
+		}
+
+		folder := viper.GetString("prompt_folder")
+		for _, key := range defaultPromptDataKeys {
+			if err := savePromptData(folder, key); err != nil {
+				return err
+			}
+		}
+		return nil
+	},
+}
+
+func savePromptData(folder, key string) error {
+	// load default prompt data
+	out, err := prompt.GetRawData(key)
+	if err != nil {
+		return err
+	}
+
+	// save out to file
+	target := path.Join(folder, key)
+	if err := os.WriteFile(target, out, 0o600); err != nil {
+		return err
+	}
+	color.Cyan("save %s to %s", key, target)
+	return nil
+}

prompt/prompt.go

@@ -28,3 +28,9 @@ func init() { //nolint:gochecknoinits
 		log.Fatal(err)
 	}
 }
+
+// GetRawData returns the raw data of the template with the given name.
+func GetRawData(name string) ([]byte, error) {
+	key := "templates/" + name
+	return templatesFS.ReadFile(key)
+}

util/template.go

@@ -6,13 +6,14 @@ import (
 	"fmt"
 	"html/template"
 	"io/fs"
+	"os"
 )
 
 // Data defines a custom type for the template data.
 type Data map[string]interface{}
 
 var (
-	templates    map[string]*template.Template
+	templates    = make(map[string]*template.Template)
 	templatesDir = "templates"
 )
 
@@ -67,9 +68,6 @@ func GetTemplateByBytes(name string, data map[string]interface{}) ([]byte, error
 // LoadTemplates loads all the templates found in the templates directory from the embedded filesystem.
 // It returns an error if reading the directory or parsing any template fails.
 func LoadTemplates(files embed.FS) error {
-	if templates == nil {
-		templates = make(map[string]*template.Template)
-	}
 	tmplFiles, err := fs.ReadDir(files, templatesDir)
 	if err != nil {
 		return err
@@ -89,3 +87,24 @@ func LoadTemplates(files embed.FS) error {
 	}
 	return nil
 }
+
+func LoadTemplatesFromDir(dir string) error {
+	tmplFiles, err := fs.ReadDir(os.DirFS(dir), ".")
+	if err != nil {
+		return err
+	}
+
+	for _, tmpl := range tmplFiles {
+		if tmpl.IsDir() {
+			continue
+		}
+
+		pt, err := template.ParseFS(os.DirFS(dir), tmpl.Name())
+		if err != nil {
+			return err
+		}
+
+		templates[tmpl.Name()] = pt
+	}
+	return nil
+}

util/template_test.go

@@ -1,7 +1,9 @@
 package util
 
 import (
+	"bytes"
 	"html/template"
+	"os"
 	"testing"
 )
 
@@ -69,3 +71,42 @@ func TestProcessTemplate(t *testing.T) {
 		t.Errorf("Unexpected output. Got: %v, Want: %v", buf.String(), expected)
 	}
 }
+
+func TestLoadTemplatesFromDir(t *testing.T) {
+	// Create a temporary directory for testing
+	tempDir := t.TempDir()
+
+	// Create a sample template file in the temporary directory
+	templateContent := "Hello, {{.Name}}!"
+	templateFile := "test.tmpl"
+	err := os.WriteFile(tempDir+"/"+templateFile, []byte(templateContent), 0o600)
+	if err != nil {
+		t.Fatalf("Failed to create test template file: %v", err)
+	}
+
+	// Load templates from the temporary directory
+	err = LoadTemplatesFromDir(tempDir)
+	if err != nil {
+		t.Fatalf("Failed to load templates from directory: %v", err)
+	}
+
+	// Check if the template was loaded correctly
+	tmpl, ok := templates[templateFile]
+	if !ok {
+		t.Fatalf("Template %s not found in loaded templates", templateFile)
+	}
+
+	// Process the loaded template
+	data := Data{"Name": "World"}
+	var buf bytes.Buffer
+	err = tmpl.Execute(&buf, data)
+	if err != nil {
+		t.Fatalf("Failed to execute loaded template: %v", err)
+	}
+
+	// Check the output
+	expected := "Hello, World!"
+	if buf.String() != expected {
+		t.Errorf("Unexpected output. Got: %v, Want: %v", buf.String(), expected)
+	}
+}