Merge pull request #50 from kardolus/variables-for-interactive-mode

README.md

@@ -21,6 +21,7 @@ Azure, featuring streaming capabilities and extensive configuration options.
 - [Getting Started](#getting-started)
 - [Configuration](#configuration)
     - [General Configuration](#general-configuration)
+      - [Variables for interactive mode](#variables-for-interactive-mode)
     - [Azure Configuration](#azure-configuration)
     - [Command-Line Autocompletion](#command-line-autocompletion)
         - [Enabling Autocompletion](#enabling-autocompletion)
@@ -192,7 +193,15 @@ Configuration variables:
 | `models_path`       | The API endpoint for accessing model information.                                                                                                      | '/v1/models'                   |
 | `auth_header`       | The header used for authorization in API requests.                                                                                                     | 'Authorization'                |
 | `auth_token_prefix` | The prefix to be added before the token in the `auth_header`.                                                                                          | 'Bearer '                      |
-| `command_prompt`    | The command prompt in interactive mode.                                                                                                                | '[YYYY-MM-DD HH:MM:SS] [Q#]: ' |
+| `command_prompt`    | The command prompt in interactive mode. Should be single-quoted.                                                                                       | '[%datetime] [Q%counter]'      |
+
+#### Variables for interactive mode:
+
+- `%date`: The current date in the format `YYYY-MM-DD`.
+- `%time`: The current time in the format `HH:MM:SS`.
+- `%datetime`: The current date and time in the format `YYYY-MM-DD HH:MM:SS`.
+- `%counter`: The total number of queries in the current session.
+- `%usage`: The usage in total tokens used (only works in query mode).
 
 The defaults can be overridden by providing your own values in the user configuration file,
 named `.chatgpt-cli/config.yaml`, located in your home directory.

cmd/chatgpt/main.go

@@ -5,7 +5,6 @@ import (
 	"fmt"
 	"io"
 	"os"
-	"strconv"
 	"strings"
 	"time"
 
@@ -209,25 +208,21 @@ func run(cmd *cobra.Command, args []string) error {
 		}
 		defer rl.Close()
 
-		prompt := func(counter string) string {
+		prompt := func(counter int) string {
 			cm := configmanager.New(config.New())
-			if len(cm.Config.CommandPrompt) != 0 {
-				return cm.Config.CommandPrompt
-			} else {
-				return fmt.Sprintf("[%s] [%s]: ", time.Now().Format("2006-01-02 15:04:05"), counter)
-			}
+			return config.FormatPrompt(cm.Config.CommandPrompt, counter, 0, time.Now())
 		}
 
 		qNum, usage := 1, 0
 		for {
 			if queryMode {
-				rl.SetPrompt(prompt(strconv.Itoa(usage)))
+				rl.SetPrompt(prompt(usage))
 			} else {
-				rl.SetPrompt(prompt(fmt.Sprintf("Q%d", qNum)))
+				rl.SetPrompt(prompt(qNum))
 			}
 
 			line, err := rl.Readline()
-			if err == readline.ErrInterrupt || err == io.EOF {
+			if errors.Is(err, readline.ErrInterrupt) || err == io.EOF {
 				fmt.Println("Bye!")
 				break
 			}

config/store.go

@@ -26,7 +26,7 @@ const (
 	openAITopP             = 1.0
 	openAIFrequencyPenalty = 0.0
 	openAIPresencePenalty  = 0.0
-	openAICommandPrompt    = ""
+	openAICommandPrompt    = "[%datetime] [Q%counter]"
 )
 
 type ConfigStore interface {

config/utils.go

@@ -0,0 +1,29 @@
+package config
+
+import (
+	"fmt"
+	"strings"
+	"time"
+)
+
+func FormatPrompt(str string, counter, usage int, now time.Time) string {
+	variables := map[string]string{
+		"%datetime": now.Format("2006-01-02 15:04:05"),
+		"%date":     now.Format("2006-01-02"),
+		"%time":     now.Format("15:04:05"),
+		"%counter":  fmt.Sprintf("%d", counter),
+		"%usage":    fmt.Sprintf("%d", usage),
+	}
+
+	// Replace placeholders in the order of longest to shortest
+	for _, key := range []string{"%datetime", "%date", "%time", "%counter", "%usage"} {
+		str = strings.ReplaceAll(str, key, variables[key])
+	}
+
+	// Ensure the last character is a space
+	if !strings.HasSuffix(str, " ") {
+		str += " "
+	}
+
+	return str
+}

config/utils_test.go

@@ -0,0 +1,81 @@
+package config_test
+
+import (
+	"fmt"
+	"github.com/kardolus/chatgpt-cli/config"
+	. "github.com/onsi/gomega"
+	"github.com/sclevine/spec"
+	"github.com/sclevine/spec/report"
+	"testing"
+	"time"
+)
+
+func TestUnitUtils(t *testing.T) {
+	spec.Run(t, "Testing the config/utils package", testUtils, spec.Report(report.Terminal{}))
+}
+
+func testUtils(t *testing.T, when spec.G, it spec.S) {
+	it.Before(func() {
+		RegisterTestingT(t)
+	})
+
+	when("FormatPrompt()", func() {
+		const (
+			counter = 1
+			usage   = 2
+		)
+
+		now := time.Now()
+
+		it("should add a trailing whitespace", func() {
+			input := "prompt"
+			expected := "prompt "
+			Expect(config.FormatPrompt(input, counter, usage, now)).To(Equal(expected))
+		})
+
+		it("should handle empty input as expected", func() {
+			input := ""
+			expected := " "
+			Expect(config.FormatPrompt(input, counter, usage, now)).To(Equal(expected))
+		})
+
+		it("should replace %date with the current date", func() {
+			currentDate := now.Format("2006-01-02")
+			input := "Today's date is %date"
+			expected := "Today's date is " + currentDate + " "
+			Expect(config.FormatPrompt(input, counter, usage, now)).To(Equal(expected))
+		})
+
+		it("should replace %time with the current time", func() {
+			currentTime := now.Format("15:04:05")
+			input := "Current time is %time"
+			expected := "Current time is " + currentTime + " "
+			Expect(config.FormatPrompt(input, counter, usage, now)).To(Equal(expected))
+		})
+
+		it("should replace %datetime with the current date and time", func() {
+			currentDatetime := now.Format("2006-01-02 15:04:05")
+			input := "Current date and time is %datetime"
+			expected := "Current date and time is " + currentDatetime + " "
+			Expect(config.FormatPrompt(input, counter, usage, now)).To(Equal(expected))
+		})
+
+		it("should replace %counter with the current counter value", func() {
+			input := "The counter is %counter"
+			expected := "The counter is 1 "
+			Expect(config.FormatPrompt(input, counter, usage, now)).To(Equal(expected))
+		})
+
+		it("should replace %usage with the current usage value", func() {
+			input := "The usage is %usage"
+			expected := "The usage is 2 "
+			Expect(config.FormatPrompt(input, counter, usage, now)).To(Equal(expected))
+		})
+
+		it("should handle complex cases correctly", func() {
+			input := "command_prompt: [%time] [Q%counter]"
+			expected := fmt.Sprintf("command_prompt: [%s] [Q%d] ", now.Format("15:04:05"), counter)
+			Expect(config.FormatPrompt(input, counter, usage, now)).To(Equal(expected))
+		})
+	})
+}