package main import ( "bufio" "bytes" "context" "embed" "fmt" "os" "runtime/debug" "strconv" "strings" gpt3 "github.com/sashabaranov/go-gpt3" "github.com/spf13/cobra" ) var LongHelp = ` Chat with ChatGPT in console. Examples: # start an interactive session chatgpt -i # ask chatgpt for a one-time response chatgpt -q "answer me this ChatGPT..." # provide context to a question or conversation chatgpt context.txt -i chatgpt context.txt -q "answer me this ChatGPT..." # read prompt from file and --write response back chatgpt convo.txt chatgpt convo.txt --write # pipe content from another program, useful for ! in vim visual mode cat convo.txt | chatgpt # inspect the predifined pretexts, which set ChatGPT's mood chatgpt -p list chatgpt -p view: # use a pretext with any of the previous modes chatgpt -p optimistic -i chatgpt -p cynic -q "Is the world going to be ok?" chatgpt -p teacher convo.txt # edit mode chatgpt -e ... # code mode chatgpt -c ... # model options (https://platform.openai.com/docs/api-reference/completions/create) chatgpt -T 4096 # set max tokens in reponse [0,4096] chatgpt -C # clean whitespace before sending chatgpt -E # echo back the prompt, useful for vim coding chatgpt --temp # set the temperature param [0.0,2.0] chatgpt --topp # set the TopP param [0.0,1.0] chatgpt --pres # set the Presence Penalty [-2.0,2.0] chatgpt --freq # set the Frequency Penalty [-2.0,2.0] ` var interactiveHelp = `starting interactive session... 'quit' to exit 'save ' to preserve 'tokens' to change the MaxToken param 'count' to change number of responses 'temp' set the temperature param [0.0,2.0] 'topp' set the TopP param [0.0,1.0] 'pres' set the Presence Penalty [-2.0,2.0] 'freq' set the Frequency Penalty [-2.0,2.0] ` //go:embed pretexts/* var predefined embed.FS var Version bool // prompt vars var Question string var Pretext string var PromptMode bool var EditMode bool var CodeMode bool var CleanPrompt bool var WriteBack bool var PromptText string // chatgpt vars var MaxTokens int var Count int var Echo bool var Temp float64 var TopP float64 var PresencePenalty float64 var FrequencyPenalty float64 var Model string // checkModel verifies that the selected mode id is one that go-gpt3 knows // about, producing an error if not. // // TODO: in future, this could probably leverage gpt3.Client.ListModels() to // support the user's fine-tuned models. func checkModel(m string) error { knownModels := []string{ gpt3.GPT3TextDavinci003, gpt3.GPT3TextDavinci002, gpt3.GPT3TextCurie001, gpt3.GPT3TextBabbage001, gpt3.GPT3TextAda001, gpt3.GPT3TextDavinci001, gpt3.GPT3DavinciInstructBeta, gpt3.GPT3Davinci, gpt3.GPT3CurieInstructBeta, gpt3.GPT3Curie, gpt3.GPT3Ada, gpt3.GPT3Babbage, } for _, v := range knownModels { if m == v { return nil } } return fmt.Errorf("unknown model '%s', expected one of: %s", m, strings.Join(knownModels, ", ")) } func GetCompletionResponse(client *gpt3.Client, ctx context.Context, question string) ([]string, error) { if CleanPrompt { question = strings.ReplaceAll(question, "\n", " ") question = strings.ReplaceAll(question, " ", " ") } // insert newline at end to prevent completion of question if !strings.HasSuffix(question, "\n") { question += "\n" } err := checkModel(Model) if err != nil { return nil, err } req := gpt3.CompletionRequest{ Model: Model, MaxTokens: MaxTokens, Prompt: question, Echo: Echo, N: Count, Temperature: float32(Temp), TopP: float32(TopP), PresencePenalty: float32(PresencePenalty), FrequencyPenalty: float32(FrequencyPenalty), } resp, err := client.CreateCompletion(ctx, req) if err != nil { return nil, err } var r []string for _, c := range resp.Choices { r = append(r, c.Text) } return r, nil } func GetEditsResponse(client *gpt3.Client, ctx context.Context, input, instruction string) ([]string, error) { if CleanPrompt { input = strings.ReplaceAll(input, "\n", " ") input = strings.ReplaceAll(input, " ", " ") } err := checkModel(Model) if err != nil { return nil, err } m := Model req := gpt3.EditsRequest{ Model: &m, Input: input, Instruction: instruction, N: Count, Temperature: float32(Temp), TopP: float32(TopP), } resp, err := client.Edits(ctx, req) if err != nil { return nil, err } var r []string for _, c := range resp.Choices { r = append(r, c.Text) } return r, nil } func GetCodeResponse(client *gpt3.Client, ctx context.Context, question string) ([]string, error) { if CleanPrompt { question = strings.ReplaceAll(question, "\n", " ") question = strings.ReplaceAll(question, " ", " ") } // insert newline at end to prevent completion of question if !strings.HasSuffix(question, "\n") { question += "\n" } req := gpt3.CompletionRequest{ Model: gpt3.CodexCodeDavinci002, MaxTokens: MaxTokens, Prompt: question, Echo: Echo, N: Count, Temperature: float32(Temp), TopP: float32(TopP), PresencePenalty: float32(PresencePenalty), FrequencyPenalty: float32(FrequencyPenalty), } resp, err := client.CreateCompletion(ctx, req) if err != nil { return nil, err } var r []string for _, c := range resp.Choices { r = append(r, c.Text) } return r, nil } func printVersion() { info, _ := debug.ReadBuildInfo() GoVersion := info.GoVersion Commit := "" BuildDate := "" dirty := false for _, s := range info.Settings { if s.Key == "vcs.revision" { Commit = s.Value } if s.Key == "vcs.time" { BuildDate = s.Value } if s.Key == "vcs.modified" { if s.Value == "true" { dirty = true } } } if dirty { Commit += "+dirty" } fmt.Printf("%s %s %s\n", Commit, BuildDate, GoVersion) } type NullWriter int func (NullWriter) Write([]byte) (int, error) { return 0, nil } func main() { apiKey := os.Getenv("CHATGPT_API_KEY") if apiKey == "" { fmt.Println("CHATGPT_API_KEY environment var is missing\nVisit https://platform.openai.com/account/api-keys to get one\n") os.Exit(1) } client := gpt3.NewClient(apiKey) rootCmd := &cobra.Command{ Use: "chatgpt [file]", Short: "Chat with ChatGPT in console.", Long: LongHelp, Run: func(cmd *cobra.Command, args []string) { if Version { printVersion() os.Exit(0) } var err error var filename string // We build up PromptText as we go, based on flags // Handle the pretext flag if Pretext != "" { files, err := predefined.ReadDir("pretexts") if err != nil { panic(err) } // list and exit if Pretext == "list" { for _, f := range files { fmt.Println(strings.TrimSuffix(f.Name(), ".txt")) } os.Exit(0) } // print pretext and exit if strings.HasPrefix(Pretext, "view:") { name := strings.TrimPrefix(Pretext, "view:") contents, err := predefined.ReadFile("pretexts/" + name + ".txt") if err != nil { fmt.Println(err) os.Exit(1) } fmt.Println(string(contents)) os.Exit(0) } // prime prompt with known pretext for _, f := range files { name := strings.TrimSuffix(f.Name(), ".txt") if name == Pretext { contents, err := predefined.ReadFile("pretexts/" + name + ".txt") if err != nil { fmt.Println(err) os.Exit(1) } PromptText = string(contents) break } } // prime prompt with custom pretext if PromptText == "" { PromptText = Pretext } } // no args, interactive, or question... read from stdin // this is mainly for replacing text in vim if len(args) == 0 && !PromptMode && Question == "" { reader := bufio.NewReader(os.Stdin) var buf bytes.Buffer for { b, err := reader.ReadByte() if err != nil { break } buf.WriteByte(b) } PromptText += buf.String() } else if len(args) == 1 { // if we have an arg, add it to the prompt filename = args[0] content, err := os.ReadFile(filename) if err != nil { fmt.Println(err) return } PromptText += string(content) } // if there is a question, it comes last in the prompt if Question != "" && !EditMode { PromptText += "\n" + Question } // interactive or file mode if PromptMode { fmt.Println(interactiveHelp) fmt.Println(PromptText) err = RunPrompt(client) } else { // empty filename (no args) prints to stdout err = RunOnce(client, filename) } if err != nil { fmt.Println(err) os.Exit(1) } }, } // setup flags rootCmd.Flags().BoolVarP(&Version, "version", "", false, "print version information") // prompt releated rootCmd.Flags().StringVarP(&Question, "question", "q", "", "ask a single question and print the response back") rootCmd.Flags().StringVarP(&Pretext, "pretext", "p", "", "pretext to add to ChatGPT input, use 'list' or 'view:' to inspect predefined, '' to use a pretext, or otherwise supply any custom text") rootCmd.Flags().BoolVarP(&PromptMode, "interactive", "i", false, "start an interactive session with ChatGPT") rootCmd.Flags().BoolVarP(&EditMode, "edit", "e", false, "request an edit with ChatGPT") rootCmd.Flags().BoolVarP(&CodeMode, "code", "c", false, "request code completion with ChatGPT") rootCmd.Flags().BoolVarP(&CleanPrompt, "clean", "x", false, "remove excess whitespace from prompt before sending") rootCmd.Flags().BoolVarP(&WriteBack, "write", "w", false, "write response to end of context file") // params related rootCmd.Flags().IntVarP(&MaxTokens, "tokens", "T", 1024, "set the MaxTokens to generate per response") rootCmd.Flags().IntVarP(&Count, "count", "C", 1, "set the number of response options to create") rootCmd.Flags().BoolVarP(&Echo, "echo", "E", false, "Echo back the prompt, useful for vim coding") rootCmd.Flags().Float64VarP(&Temp, "temp", "", 1.0, "set the temperature parameter") rootCmd.Flags().Float64VarP(&TopP, "topp", "", 1.0, "set the TopP parameter") rootCmd.Flags().Float64VarP(&PresencePenalty, "pres", "", 0.0, "set the Presence Penalty parameter") rootCmd.Flags().Float64VarP(&FrequencyPenalty, "freq", "", 0.0, "set the Frequency Penalty parameter") rootCmd.Flags().StringVarP(&Model, "model", "m", gpt3.GPT3TextDavinci003, "select the model to use with -q or -e") // run the command rootCmd.Execute() } func RunPrompt(client *gpt3.Client) error { ctx := context.Background() scanner := bufio.NewScanner(os.Stdin) quit := false for !quit { fmt.Print("> ") if !scanner.Scan() { break } question := scanner.Text() parts := strings.Fields(question) // look for commands switch parts[0] { case "quit", "q", "exit": quit = true continue case "save": name := parts[1] fmt.Printf("saving session to %s\n", name) err := os.WriteFile(name, []byte(PromptText), 0644) if err != nil { fmt.Println(err) } continue case "tokens": if len(parts) == 1 { fmt.Println("tokens is set to", MaxTokens) continue } c, err := strconv.Atoi(parts[1]) if err != nil { fmt.Println(err) continue } MaxTokens = c fmt.Println("tokens is now", MaxTokens) continue case "count": if len(parts) == 1 { fmt.Println("count is set to", Count) continue } c, err := strconv.Atoi(parts[1]) if err != nil { fmt.Println(err) continue } Count = c fmt.Println("count is now", Count) continue case "temp": if len(parts) == 1 { fmt.Println("temp is set to", Temp) continue } f, err := strconv.ParseFloat(parts[1], 64) if err != nil { fmt.Println(err) continue } Temp = f fmt.Println("temp is now", Temp) case "topp": if len(parts) == 1 { fmt.Println("topp is set to", TopP) continue } f, err := strconv.ParseFloat(parts[1], 64) if err != nil { fmt.Println(err) continue } TopP = f fmt.Println("topp is now", TopP) case "pres": if len(parts) == 1 { fmt.Println("pres is set to", PresencePenalty) continue } f, err := strconv.ParseFloat(parts[1], 64) if err != nil { fmt.Println(err) continue } PresencePenalty = f fmt.Println("pres is now", PresencePenalty) case "freq": if len(parts) == 1 { fmt.Println("freq is set to", FrequencyPenalty) continue } f, err := strconv.ParseFloat(parts[1], 64) if err != nil { fmt.Println(err) continue } FrequencyPenalty = f fmt.Println("freq is now", FrequencyPenalty) default: // add the question to the existing prompt text, to keep context PromptText += "\n> " + question var R []string var err error if CodeMode { R, err = GetCodeResponse(client, ctx, PromptText) } else if EditMode { R, err = GetEditsResponse(client, ctx, PromptText, Question) } else { R, err = GetCompletionResponse(client, ctx, PromptText) } if err != nil { return err } final := "" if len(R) == 1 { final = R[0] } else { for i, r := range R { final += fmt.Sprintf("[%d]: %s\n\n", i, r) } fmt.Println(final) ok := false pos := 0 for !ok { fmt.Print("> ") if !scanner.Scan() { break } ans := scanner.Text() pos, err = strconv.Atoi(ans) if err != nil { fmt.Println(err) continue } if pos < 0 || pos >= Count { fmt.Println("choice must be between 0 and", Count-1) continue } ok = true } final = R[pos] } // we add response to the prompt, this is how ChatGPT sessions keep context PromptText += "\n" + strings.TrimSpace(final) // print the latest portion of the conversation fmt.Println(final + "\n") } } return nil } func RunOnce(client *gpt3.Client, filename string) error { ctx := context.Background() var R []string var err error if CodeMode { R, err = GetCodeResponse(client, ctx, PromptText) } else if EditMode { R, err = GetEditsResponse(client, ctx, PromptText, Question) } else { R, err = GetCompletionResponse(client, ctx, PromptText) } if err != nil { return err } final := "" if len(R) == 1 { final = R[0] } else { for i, r := range R { final += fmt.Sprintf("[%d]: %s\n\n", i, r) } } if filename == "" || !WriteBack { fmt.Println(final) } else { err = AppendToFile(filename, final) if err != nil { return err } } return nil } // AppendToFile provides a function to append data to an existing file, // creating it if it doesn't exist func AppendToFile(filename string, data string) error { // Open the file in append mode file, err := os.OpenFile(filename, os.O_APPEND|os.O_CREATE|os.O_WRONLY, 0644) if err != nil { return err } // Append the data to the file _, err = file.WriteString(data) if err != nil { return err } return file.Close() }