package main
import (
"bufio"
"bytes"
"fmt"
"os"
"runtime/debug"
"github.com/sashabaranov/go-openai"
"github.com/spf13/cobra"
)
var LongHelp = `
Work with ChatGPT in console .
Examples :
# start an interactive session with gpt - 3.5 or gpt - 4
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
# set the directory for custom prompts , defaults to "prompts"
chatgpt - P custom - prompts - p my - prompt - i
# inspect the predifined prompts , which set ChatGPT ' s mood
chatgpt - p list
chatgpt - p view : < name >
# use a prompts 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 ]
# change model selection , available models are listed here :
# https : //pkg.go.dev/github.com/sashabaranov/go-openai#Client.ListModels
chatgpt - m text - davinci - 003 # set the model to text - davinci - 003 ( the default )
chatgpt - m text - ada - 001 # set the model to text - ada - 001
`
var Version bool
// prompt vars
var Question string
var Prompt string
var PromptDir 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
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 ( "OPENAI_API_KEY" )
if apiKey == "" {
fmt . Println ( "OPENAI_API_KEY environment var is missing\nVisit https://platform.openai.com/account/api-keys to get one\n" )
os . Exit ( 1 )
}
if PromptDir == "" {
if v := os . Getenv ( "CHATGPT_PROMPT_DIR" ) ; v != "" {
PromptDir = v
}
}
client := openai . 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 prompt flag
if Prompt != "" {
p , err := handlePrompt ( Prompt )
if err != nil {
fmt . Println ( err )
os . Exit ( 1 )
}
if p == "" {
os . Exit ( 0 )
}
// prime prompt with custom pretext
PromptText = p
}
// no args, or interactive... read from stdin
// this is mainly for replacing text in vim
fi , _ := os . Stdin . Stat ( )
inputIsFromPipe := ( fi . Mode ( ) & os . ModeCharDevice ) == 0
if len ( args ) == 0 && ! PromptMode && inputIsFromPipe {
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 {
err = RunInteractive ( 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 ( & Prompt , "prompt" , "p" , "" , "prompt to add to ChatGPT input, use 'list' or 'view:<name>' to inspect predefined, '<name>' to use a prompt, or otherwise supply any custom text" )
rootCmd . Flags ( ) . StringVarP ( & PromptDir , "prompt-dir" , "P" , "prompts" , "directory containing custom prompts, if not set the embedded defaults are used" )
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 ( & 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" , "" , 0.7 , "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" , openai . GPT3TextDavinci003 , "select the model to use with -q or -e" )
// run the command
rootCmd . Execute ( )
}