From 46e1adaec16905951988e88701728b80cc1f28c5 Mon Sep 17 00:00:00 2001 From: Peter Bakkum Date: Wed, 1 Mar 2023 00:03:40 -0800 Subject: [PATCH] Rewrite command improvements and fixes --- butterfish/butterfish.go | 32 +++++++++++++++++----- butterfish/commands.go | 57 +++++++++++++++++++++++++++++++++++----- cmd/butterfish/main.go | 1 + go.mod | 1 + go.sum | 5 +++- util/util.go | 13 ++++++--- 6 files changed, 91 insertions(+), 18 deletions(-) diff --git a/butterfish/butterfish.go b/butterfish/butterfish.go index 92b9788..a883f84 100644 --- a/butterfish/butterfish.go +++ b/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)), diff --git a/butterfish/commands.go b/butterfish/commands.go index 143ee7d..023e2e0 100644 --- a/butterfish/commands.go +++ b/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 }) diff --git a/cmd/butterfish/main.go b/cmd/butterfish/main.go index 50f8b63..a5a6cd5 100644 --- a/cmd/butterfish/main.go +++ b/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) diff --git a/go.mod b/go.mod index 8f1bc5f..da1ce55 100644 --- a/go.mod +++ b/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 diff --git a/go.sum b/go.sum index 4e662f6..bc623d7 100644 --- a/go.sum +++ b/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= diff --git a/util/util.go b/util/util.go index 13979cd..4efe892 100644 --- a/util/util.go +++ b/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