add switch to disable color and fix tests

Author: jaffee

pkg/aicli/cmd.go

@@ -33,6 +33,7 @@ type Cmd struct {
 	Verbose      bool    `help:"Enables debug output."`
 	ContextLimit int     `help:"Maximum number of bytes of context to keep. Earlier parts of the conversation are discarded."`
 	EnableAWS    bool    `help:"Enable AWS Bedrock as an AI option. Disabled by default because it slows startup time."`
+	Color        bool    `help:"Enable colored output"`
 
 	rl *readline.Instance
 
@@ -55,6 +56,7 @@ func NewCmd() *Cmd {
 		Model:        "gpt-3.5-turbo",
 		Temperature:  0.7,
 		ContextLimit: 10000, // 10,000 bytes ~2000 tokens
+		Color:        true,
 
 		messages: []Message{},
 
@@ -82,15 +84,7 @@ func (cmd *Cmd) Run() error {
 	}
 
 	var err error
-	cmd.rl, err = readline.NewEx(&readline.Config{
-		Prompt:       "> ",
-		HistoryFile:  cmd.getHistoryFilePath(),
-		HistoryLimit: 1000000,
-
-		Stdin:  cmd.stdin,
-		Stdout: &colorWriter{w: cmd.stdout, color: colorCyan}, // Color the user input
-		Stderr: cmd.stderr,
-	})
+	cmd.rl, err = readline.NewEx(cmd.getReadlineConfig())
 	if err != nil {
 		return err
 	}
@@ -138,7 +132,7 @@ func (cmd *Cmd) sendMessages() error {
 		Temperature: cmd.Temperature,
 		Messages:    cmd.messagesWithinLimit(),
 	}
-	msg, err := cmd.client().GenerateStream(req, &colorWriter{w: cmd.stdout, color: colorYellow})
+	msg, err := cmd.client().GenerateStream(req, &colorWriter{w: cmd.stdout, color: cmd.getColor(colorYellow)})
 	if err != nil {
 		return err
 	}
@@ -286,13 +280,16 @@ func (cmd *Cmd) readMulti(until string) (string, error) {
 func (cmd *Cmd) Set(param, value string) error {
 	switch param {
 	case "ai":
+		if _, ok := cmd.ais[value]; !ok {
+			return errors.Errorf("unknown ai '%s'", value)
+		}
 		cmd.AI = value
 	case "model":
 		cmd.Model = value
-	case "temperature":
+	case "temp", "temperature":
 		temp, err := strconv.ParseFloat(value, 64)
 		if err != nil {
-			return errors.Wrap(err, "parsing temp value")
+			return errors.Wrapf(err, "parsing '%s' to float", value)
 		}
 		cmd.Temperature = temp
 	case "verbose":
@@ -304,12 +301,27 @@ func (cmd *Cmd) Set(param, value string) error {
 		default:
 			return errors.Errorf("could not parse '%s' as bool", value)
 		}
-	case "context-limit":
+	case "context":
 		lim, err := strconv.Atoi(value)
 		if err != nil {
 			return errors.Wrapf(err, "parsing '%s' to int", value)
 		}
 		cmd.ContextLimit = lim
+	case "color":
+		switch strings.ToLower(value) {
+		case "1", "true", "yes", "on":
+			cmd.Color = true
+		case "0", "false", "no", "off":
+			cmd.Color = false
+		default:
+			return errors.Errorf("could not parse '%s' as bool", value)
+		}
+		// Update readline config with new color setting
+		var err error
+		cmd.rl, err = readline.NewEx(cmd.getReadlineConfig())
+		if err != nil {
+			return errors.Wrap(err, "updating readline config")
+		}
 	default:
 		return errors.Errorf("unknown parameter '%s'", param)
 	}
@@ -349,9 +361,9 @@ func (cmd *Cmd) appendMessage(msg Message) {
 func (cmd *Cmd) printMessages() {
 	for _, msg := range cmd.messages {
 		if msg.Role() == "user" {
-			cmd.out(colorCyan+"%9s: %s"+colorReset, msg.Role(), msg.Content())
+			cmd.out(cmd.getColor(colorCyan)+"%9s: %s"+cmd.getColor(colorReset), msg.Role(), msg.Content())
 		} else {
-			cmd.out(colorYellow+"%9s: %s"+colorReset, msg.Role(), msg.Content())
+			cmd.out(cmd.getColor(colorYellow)+"%9s: %s"+cmd.getColor(colorReset), msg.Role(), msg.Content())
 		}
 	}
 }
@@ -367,16 +379,16 @@ func (cmd *Cmd) out(format string, a ...any) {
 }
 
 func (cmd *Cmd) sysOut(format string, a ...any) {
-	fmt.Fprintf(cmd.stdout, colorBlue+format+colorReset+"\n", a...)
+	fmt.Fprintf(cmd.stdout, cmd.getColor(colorBlue)+format+cmd.getColor(colorReset)+"\n", a...)
 }
 
 func (cmd *Cmd) err(err error) {
-	fmt.Fprintf(cmd.stderr, colorRed+"%s"+colorReset+"\n", err.Error())
+	fmt.Fprintf(cmd.stderr, cmd.getColor(colorRed)+"%s"+cmd.getColor(colorReset)+"\n", err.Error())
 }
 
 // errOut wraps the error and writes it to the user on stderr.
 func (cmd *Cmd) errOut(err error, format string, a ...any) {
-	fmt.Fprintf(cmd.stderr, colorRed+"%s: %v"+colorReset+"\n", fmt.Sprintf(format, a...), err.Error())
+	fmt.Fprintf(cmd.stderr, cmd.getColor(colorRed)+"%s: %v"+cmd.getColor(colorReset)+"\n", fmt.Sprintf(format, a...), err.Error())
 }
 
 // checkConfig ensures the command configuration is valid before proceeding.
@@ -393,6 +405,7 @@ func (cmd *Cmd) printConfig() {
 	fmt.Fprintf(cmd.stderr, "Temperature: %f\n", cmd.Temperature)
 	fmt.Fprintf(cmd.stderr, "Verbose: %v\n", cmd.Verbose)
 	fmt.Fprintf(cmd.stderr, "ContextLimit: %d\n", cmd.ContextLimit)
+	fmt.Fprintf(cmd.stderr, "Color: %v\n", cmd.Color)
 }
 
 func (cmd *Cmd) setupConfigDir() error {
@@ -444,13 +457,27 @@ func (cmd *Cmd) readConfigFile() error {
 	return nil
 }
 
+func (cmd *Cmd) getReadlineConfig() *readline.Config {
+	return &readline.Config{
+		Prompt:       "> ",
+		HistoryFile:  cmd.getHistoryFilePath(),
+		HistoryLimit: 1000000,
+		Stdin:        cmd.stdin,
+		Stdout:       &colorWriter{w: cmd.stdout, color: cmd.getColor(colorCyan)},
+		Stderr:       cmd.stderr,
+	}
+}
+
 // colorWriter wraps an io.Writer and adds color codes
 type colorWriter struct {
 	w     io.Writer
 	color string
 }
 
 func (cw *colorWriter) Write(p []byte) (n int, err error) {
+	if cw.color == "" {
+		return cw.w.Write(p)
+	}
 	// Write the color code, then the content, then reset
 	if _, err := cw.w.Write([]byte(cw.color)); err != nil {
 		return 0, err
@@ -462,3 +489,10 @@ func (cw *colorWriter) Write(p []byte) (n int, err error) {
 	_, err = cw.w.Write([]byte(colorReset))
 	return n, err
 }
+
+func (cmd *Cmd) getColor(col string) string {
+	if !cmd.Color {
+		return ""
+	}
+	return col
+}

pkg/aicli/cmd_test.go

@@ -1,6 +1,7 @@
 package aicli
 
 import (
+	"fmt"
 	"io"
 	"testing"
 	"time"
@@ -11,6 +12,7 @@ import (
 func TestCmd(t *testing.T) {
 	cmd := NewCmd()
 	cmd.AI = "echo"
+	cmd.Color = false
 	cmd.AddAI("echo", &Echo{})
 	stdinr, stdinw := io.Pipe()
 	stdout, stdoutw := io.Pipe()
@@ -25,6 +27,9 @@ func TestCmd(t *testing.T) {
 	var runErr error
 	go func() {
 		runErr = cmd.Run()
+		if runErr != nil {
+			fmt.Printf("error: %v\n", runErr)
+		}
 		close(done)
 	}()
 
@@ -42,7 +47,7 @@ func TestCmd(t *testing.T) {
 	write(t, stdinw, []byte("\\reset\n"))
 	require.NoError(t, runErr)
 	write(t, stdinw, []byte("\\config\n"))
-	expect(t, stderr, []byte("AI: echo\nModel: gpt-3.5-turbo\nTemperature: 0.700000\nVerbose: false\nContextLimit: 10000\n"))
+	expect(t, stderr, []byte("AI: echo\nModel: gpt-3.5-turbo\nTemperature: 0.700000\nVerbose: false\nContextLimit: 10000\nColor: false\n"))
 	write(t, stdinw, []byte("\\reset\n"))
 	write(t, stdinw, []byte("\\file ./testdata/myfile\n"))
 	write(t, stdinw, []byte("\\messages\n"))