Rewrite command improvements and fixes

butterfish/butterfish.go

@@ -67,10 +67,10 @@ type ButterfishCtx struct {
 }
 
 type ColorScheme struct {
-	Foreground string
+	Foreground string // neutral foreground color
 	Background string
-	Error      string
-	Color1     string
+	Error      string // should be reddish
+	Color1     string // should be greenish
 	Color2     string
 	Color3     string
 	Color4     string
@@ -85,9 +85,9 @@ var GruvboxDark = ColorScheme{
 	Foreground: "#ebdbb2",
 	Background: "#282828",
 	Error:      "#fb4934", // red
-	Color1:     "#bb8b26", // green
+	Color1:     "#b8bb26", // green
 	Color2:     "#fabd2f", // yellow
-	Color3:     "#458588", // blue
+	Color3:     "#83a598", // blue
 	Color4:     "#d3869b", // magenta
 	Color5:     "#8ec07c", // cyan
 	Color6:     "#fe8019", // orange
@@ -237,6 +237,10 @@ func (this *ButterfishCtx) StylePrintf(style lipgloss.Style, format string, a ..
 	this.Out.Write([]byte(str))
 }
 
+func (this *ButterfishCtx) StyleSprintf(style lipgloss.Style, format string, a ...any) string {
+	return util.MultilineLipglossRender(style, fmt.Sprintf(format, a...))
+}
+
 func (this *ButterfishCtx) Printf(format string, a ...any) {
 	this.StylePrintf(this.Config.Styles.Foreground, format, a...)
 }
@@ -287,6 +291,7 @@ func (this *ButterfishCtx) printError(err error, prefix ...string) {
 type styles struct {
 	Question   lipgloss.Style
 	Answer     lipgloss.Style
+	Go         lipgloss.Style
 	Summarize  lipgloss.Style
 	Highlight  lipgloss.Style
 	Prompt     lipgloss.Style
@@ -295,10 +300,23 @@ type styles struct {
 	Grey       lipgloss.Style
 }
 
+func (this *styles) PrintTestColors() {
+	fmt.Println("Question: ", this.Question.Render("Question"))
+	fmt.Println("Answer: ", this.Answer.Render("Answer"))
+	fmt.Println("Go: ", this.Go.Render("Go"))
+	fmt.Println("Summarize: ", this.Summarize.Render("Summarize"))
+	fmt.Println("Highlight: ", this.Highlight.Render("Highlight"))
+	fmt.Println("Prompt: ", this.Prompt.Render("Prompt"))
+	fmt.Println("Error: ", this.Error.Render("Error"))
+	fmt.Println("Foreground: ", this.Foreground.Render("Foreground"))
+	fmt.Println("Grey: ", this.Grey.Render("Grey"))
+}
+
 func ColorSchemeToStyles(colorScheme *ColorScheme) *styles {
 	return &styles{
-		Question:   lipgloss.NewStyle().Foreground(lipgloss.Color(colorScheme.Color3)),
-		Answer:     lipgloss.NewStyle().Foreground(lipgloss.Color(colorScheme.Color1)),
+		Question:   lipgloss.NewStyle().Foreground(lipgloss.Color(colorScheme.Color4)),
+		Answer:     lipgloss.NewStyle().Foreground(lipgloss.Color(colorScheme.Color2)),
+		Go:         lipgloss.NewStyle().Foreground(lipgloss.Color(colorScheme.Color5)),
 		Highlight:  lipgloss.NewStyle().Foreground(lipgloss.Color(colorScheme.Color2)),
 		Summarize:  lipgloss.NewStyle().Foreground(lipgloss.Color(colorScheme.Color2)),
 		Prompt:     lipgloss.NewStyle().Foreground(lipgloss.Color(colorScheme.Color4)),

butterfish/commands.go

@@ -14,6 +14,7 @@ import (
 	"github.com/alecthomas/kong"
 	"github.com/bakks/butterfish/prompt"
 	"github.com/bakks/butterfish/util"
+	"github.com/sergi/go-diff/diffmatchpatch"
 	"github.com/spf13/afero"
 )
 
@@ -68,7 +69,7 @@ type CliCommandConfig struct {
 		Prompt      string  `arg:"" help:"Instruction to the model on how to rewrite."`
 		Inputfile   string  `short:"i" help:"Source file for content to rewrite. If not set then there must be piped input."`
 		Outputfile  string  `short:"o" help:"File to write the rewritten output to."`
-		Inplace     bool    `short:"I" help:"Rewrite the input file in place, cannot be set at the same time as the outputfile flag."`
+		Inplace     bool    `short:"I" help:"Rewrite the input file in place. This is potentially destructive, use with caution! Cannot be set at the same time as the outputfile flag."`
 		Model       string  `short:"m" default:"code-davinci-edit-001" help:"GPT model to use for editing. At compile time this should be either 'code-davinci-edit-001' or 'text-davinci-edit-001'."`
 		Temperature float32 `short:"T" default:"0.6" help:"Temperature to use for the prompt, higher temperature indicates more freedom/randomness when generating each token."`
 		ChunkSize   int     `short:"c" default:"4000" help:"Number of bytes to rewrite at a time if the file must be split up."`
@@ -429,6 +430,25 @@ func (this *ButterfishCtx) Prompt(prompt string, model string, maxTokens int, te
 	return err
 }
 
+func (this *ButterfishCtx) diffStrings(a, b string) string {
+	dmp := diffmatchpatch.New()
+	diffs := dmp.DiffMain(a, b, false)
+
+	strBuilder := strings.Builder{}
+	for _, diff := range diffs {
+		switch diff.Type {
+		case diffmatchpatch.DiffInsert:
+			strBuilder.WriteString(
+				this.StyleSprintf(this.Config.Styles.Go, "%s", diff.Text))
+		case diffmatchpatch.DiffEqual:
+			strBuilder.WriteString(
+				this.StyleSprintf(this.Config.Styles.Foreground, "%s", diff.Text))
+		}
+	}
+
+	return strBuilder.String()
+}
+
 func (this *ButterfishCtx) rewriteCommand(
 	prompt, model string,
 	inputFilePath, outputFilePath string,
@@ -458,6 +478,7 @@ func (this *ButterfishCtx) rewriteCommand(
 
 	var inputReader io.Reader
 	var outputWriter io.Writer
+	doStyling := false
 
 	inputReader = this.getPipedStdinReader()
 	if inputReader != nil && inputFilePath != "" {
@@ -478,14 +499,31 @@ func (this *ButterfishCtx) rewriteCommand(
 
 	if outputFilePath != "" {
 		// open output file for writing
-		outputFile, err := os.Create(outputFilePath)
-		defer outputFile.Close()
-		if err != nil {
-			return err
+		// if the output file is the same as the input file then we write to a
+		// temporary file and then rename it to the input file
+		if outputFilePath == inputFilePath {
+			tempFile, err := ioutil.TempFile("", "butterfish")
+			if err != nil {
+				return err
+			}
+			defer func() {
+				tempFile.Close()
+				// move the temp file to the input file
+				os.Rename(tempFile.Name(), inputFilePath)
+			}()
+
+			outputWriter = tempFile
+		} else {
+			outputFile, err := os.Create(outputFilePath)
+			if err != nil {
+				return err
+			}
+			defer outputFile.Close()
+			outputWriter = outputFile
 		}
-		outputWriter = outputFile
 	} else {
-		outputWriter = util.NewStyledWriter(this.Out, this.Config.Styles.Answer)
+		doStyling = true
+		outputWriter = this.Out
 	}
 
 	return util.ChunkFromReader(inputReader, chunkSize, maxChunks, func(i int, chunk []byte) error {
@@ -499,6 +537,11 @@ func (this *ButterfishCtx) rewriteCommand(
 			edited = edited[:len(edited)-1]
 		}
 
+		if doStyling {
+			// do diffing and styling
+			edited = this.diffStrings(string(chunk), edited)
+		}
+
 		_, err = outputWriter.Write([]byte(edited))
 		return err
 	})

cmd/butterfish/main.go

@@ -167,6 +167,7 @@ func main() {
 			fmt.Fprintf(errorWriter, err.Error())
 			os.Exit(3)
 		}
+		//butterfishCtx.Config.Styles.PrintTestColors()
 
 		err = butterfishCtx.ExecCommand(parsedCmd, &cli.CliCommandConfig)
 

go.mod

@@ -14,6 +14,7 @@ require (
 	github.com/joho/godotenv v1.4.0
 	github.com/mitchellh/go-homedir v1.1.0
 	github.com/muesli/reflow v0.3.0
+	github.com/sergi/go-diff v1.3.1
 	github.com/spf13/afero v1.9.3
 	github.com/stretchr/testify v1.8.1
 	golang.org/x/term v0.5.0

go.sum

@@ -207,6 +207,8 @@ github.com/rogpeppe/go-internal v1.3.0/go.mod h1:M8bDsm7K2OlrFYOpmOWEs/qY81heoFR
 github.com/sahilm/fuzzy v0.1.0/go.mod h1:VFvziUEIMCrT6A6tw2RFIXPXXmzXbOsSHF0DOI8ZK9Y=
 github.com/sclevine/spec v1.2.0/go.mod h1:W4J29eT/Kzv7/b9IWLB055Z+qvVC9vt0Arko24q7p+U=
 github.com/sclevine/spec v1.4.0/go.mod h1:LvpgJaFyvQzRvc1kaDs0bulYwzC70PbiYjC4QnFHkOM=
+github.com/sergi/go-diff v1.3.1 h1:xkr+Oxo4BOQKmkn/B9eMK0g5Kg/983T9DqqPHwYqD+8=
+github.com/sergi/go-diff v1.3.1/go.mod h1:aMJSSKb2lpPvRNec0+w3fl7LP9IOFzdc9Pa4NFbPK1I=
 github.com/spf13/afero v1.9.3 h1:41FoI0fD7OR7mGcKE/aOiLkGreyf8ifIOQmJANWogMk=
 github.com/spf13/afero v1.9.3/go.mod h1:iUV7ddyEEZPO5gA3zD4fJt6iStLlL+Lg4m2cihcDf8Y=
 github.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME=
@@ -533,8 +535,9 @@ google.golang.org/protobuf v1.26.0/go.mod h1:9q0QmTI4eRPtz6boOQmLYwt+qCgq0jsYwAQ
 google.golang.org/protobuf v1.27.1 h1:SnqbnDw1V7RiZcXPx5MEeqPv2s79L9i7BJUlG/+RurQ=
 google.golang.org/protobuf v1.27.1/go.mod h1:9q0QmTI4eRPtz6boOQmLYwt+qCgq0jsYwAQnmE0givc=
 gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=
-gopkg.in/check.v1 v1.0.0-20180628173108-788fd7840127 h1:qIbj1fsPNlZgppZ+VLlY7N33q108Sa+fhmuc+sWQYwY=
 gopkg.in/check.v1 v1.0.0-20180628173108-788fd7840127/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=
+gopkg.in/check.v1 v1.0.0-20190902080502-41f04d3bba15 h1:YR8cESwS4TdDjEe65xsg0ogRM/Nc3DYOhEAlW+xobZo=
+gopkg.in/check.v1 v1.0.0-20190902080502-41f04d3bba15/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=
 gopkg.in/errgo.v2 v2.1.0/go.mod h1:hNsd1EY+bozCKY1Ytp96fpM3vjJbqLJn88ws8XvfDNI=
 gopkg.in/fsnotify.v1 v1.4.7/go.mod h1:Tz8NjZHkW78fSQdbUxIjBTcgA1z1m8ZHf0WmKUhAMys=
 gopkg.in/tomb.v1 v1.0.0-20141024135613-dd632973f1e7/go.mod h1:dt/ZhP58zS4L8KSrWDmTeBkI65Dw0HsyUHuEVlX15mw=

util/util.go

@@ -6,6 +6,7 @@ import (
 	"os"
 	"path/filepath"
 	"strings"
+	"unicode"
 
 	"github.com/charmbracelet/lipgloss"
 	"github.com/spf13/afero"
@@ -171,9 +172,10 @@ func (this *CacheWriter) GetLastN(n int) []byte {
 // and filters out the special token "NOOP". This is specially handled -
 // we seem to get "NO" as a separate token from GPT.
 type StyledWriter struct {
-	Writer io.Writer
-	Style  lipgloss.Style
-	cache  []byte
+	Writer    io.Writer
+	Style     lipgloss.Style
+	cache     []byte
+	seenInput bool
 }
 
 // Lipgloss is a little tricky - if you render a string with newlines it
@@ -200,6 +202,11 @@ func MultilineLipglossRender(style lipgloss.Style, str string) string {
 // This is a bit insane but it's a dumb way to filter out NOOP split into
 // two tokens, should probably be rewritten
 func (this *StyledWriter) Write(input []byte) (int, error) {
+	if !this.seenInput && unicode.IsSpace(rune(input[0])) {
+		return len(input), nil
+	}
+	this.seenInput = true
+
 	if string(input) == "NOOP" {
 		// This doesn't seem to actually happen since it gets split into two
 		// tokens? but let's code defensively