From 4958ad75b9cf262752f44a007a676855202b7c5a Mon Sep 17 00:00:00 2001 From: Mike Brown Date: Tue, 14 Nov 2023 22:59:11 +0000 Subject: [PATCH] Add support for querying billing/usage --- .circleci/config.yml | 2 - Makefile | 1 + cmd/gptcli/help.txt | 2 +- cmd/gptcli/main.go | 118 +++++++++++++++++++++++++++++++++++++--- cmd/gptcli/main_test.go | 67 +++++++++++++++++++++++ go.mod | 7 +-- go.sum | 7 +-- internal/interfaces.go | 2 + 8 files changed, 186 insertions(+), 20 deletions(-) diff --git a/.circleci/config.yml b/.circleci/config.yml index 9078201..cd5ede8 100644 --- a/.circleci/config.yml +++ b/.circleci/config.yml @@ -35,12 +35,10 @@ jobs: command: | mkdir -p $HOME/go/bin go install github.com/golang/mock/mockgen@latest - ls $HOME/go/bin - run: name: Run tests command: | PATH=$PATH:$HOME/go/bin - ls $HOME/go/bin make unit-tests.xml - store_test_results: path: unit-tests.xml diff --git a/Makefile b/Makefile index bed7ac1..a9c4dac 100644 --- a/Makefile +++ b/Makefile @@ -36,6 +36,7 @@ clean: deps: rm -rf go.mod go.sum vendor go mod init github.com/mikeb26/gptcli + go mod edit -replace=github.com/sashabaranov/go-openai=github.com/mikeb26/sashabaranov-go-openai@v1.17.6.mb1 GOPROXY=direct go mod tidy go mod vendor diff --git a/cmd/gptcli/help.txt b/cmd/gptcli/help.txt index 4e4410c..fcf0e38 100644 --- a/cmd/gptcli/help.txt +++ b/cmd/gptcli/help.txt @@ -11,4 +11,4 @@ Available Commands: ls List available threads(conversations) thread Switch to a previously created thread exit Exit gptcli - + billing Show billing/daily usage information diff --git a/cmd/gptcli/main.go b/cmd/gptcli/main.go index b46e1d3..f43b255 100644 --- a/cmd/gptcli/main.go +++ b/cmd/gptcli/main.go @@ -31,6 +31,7 @@ import ( const ( CommandName = "gptcli" KeyFile = ".openai.key" + SessionFile = ".openai.session" ThreadsDir = "threads" CodeBlockDelim = "```" CodeBlockDelimNewline = "```\n" @@ -54,6 +55,7 @@ var subCommandTab = map[string]func(ctx context.Context, "delete": deleteThreadMain, "exit": exitMain, "quit": exitMain, + "billing": billingMain, } type GptCliThread struct { @@ -68,11 +70,13 @@ type GptCliThread struct { type GptCliContext struct { client internal.OpenAIClient + sessClient internal.OpenAIClient input *bufio.Reader needConfig bool curThreadNum int totThreads int threads []*GptCliThread + haveSess bool } func NewGptCliContext() *GptCliContext { @@ -85,6 +89,13 @@ func NewGptCliContext() *GptCliContext { } else { clientLocal = openai.NewClient(keyText) } + var sessClientLocal *openai.Client + haveSessLocal := false + sessText, err := loadSess() + if err == nil { + sessClientLocal = openai.NewClient(sessText) + haveSessLocal = true + } return &GptCliContext{ client: clientLocal, @@ -93,6 +104,8 @@ func NewGptCliContext() *GptCliContext { curThreadNum: 0, totThreads: 0, threads: make([]*GptCliThread, 0), + haveSess: haveSessLocal, + sessClient: sessClientLocal, } } @@ -165,6 +178,43 @@ func exitMain(ctx context.Context, gptCliCtx *GptCliContext, return nil } +func billingMain(ctx context.Context, gptCliCtx *GptCliContext, + args []string) error { + + if !gptCliCtx.haveSess { + return fmt.Errorf("A session key must first be configured to use the billing feature. try 'config'") + } + endDate := time.Now().Add(24 * time.Hour) + startDate := endDate.Add(-(30 * 24 * time.Hour)) + resp, err := gptCliCtx.sessClient.GetBillingUsage(ctx, startDate, endDate) + if err != nil { + return err + } + + fmt.Printf("Usage from %v - %v:\n", startDate.Format(time.DateOnly), + endDate.Format(time.DateOnly)) + + var printedDate bool + for _, d := range resp.DailyCosts { + printedDate = false + for _, li := range d.LineItems { + if li.Cost == 0 { + continue + } + + if !printedDate { + fmt.Printf("%v:\n", d.Time.Format(time.DateOnly)) + printedDate = true + } + fmt.Printf("\t%v: $%.2f\n", li.Name, li.Cost*0.01) + } + } + + fmt.Printf("\nTotal: $%.2f\n", resp.TotalUsage*0.01) + + return nil +} + //go:embed version.txt var versionText string @@ -194,11 +244,14 @@ func upgradeMain(ctx context.Context, gptCliCtx *GptCliContext, args []string) e fmt.Printf("A new version of gptcli is available (%v). Upgrade? (Y/N) [Y]: ", latestVer) shouldUpgrade, err := gptCliCtx.input.ReadString('\n') - - shouldUpgrade = strings.ToUpper(strings.TrimSpace(shouldUpgrade)) if err != nil { return err } + + shouldUpgrade = strings.ToUpper(strings.TrimSpace(shouldUpgrade)) + if len(shouldUpgrade) == 0 { + shouldUpgrade = "Y" + } if shouldUpgrade[0] != 'Y' { return nil } @@ -345,18 +398,18 @@ func lsThreadsMain(ctx context.Context, gptCliCtx *GptCliContext, return nil } - rowFmt := "| %8v | %-16v | %-16v | %-16v | %-16v\n" - rowSpacer := "--------------------------------------------------------------------------------------------\n" + rowFmt := "| %8v | %17v | %17v | %17v | %-17v\n" + rowSpacer := "-------------------------------------------------------------------------------------------------\n" fmt.Printf(rowSpacer) fmt.Printf(rowFmt, "Thread#", "Last Accessed", "Last Modified", "Created", "Name") fmt.Printf(rowSpacer) for idx, t := range gptCliCtx.threads { - cTime := t.CreateTime.Format("1/2/2006 3:04pm") - aTime := t.AccessTime.Format("1/2/2006 3:04pm") - mTime := t.ModTime.Format("1/2/2006 3:04pm") - today := time.Now().UTC().Truncate(24 * time.Hour).Format("1/2/2006") + cTime := t.CreateTime.Format("1/02/2006 3:04pm") + aTime := t.AccessTime.Format("1/02/2006 3:04pm") + mTime := t.ModTime.Format("1/02/2006 3:04pm") + today := time.Now().UTC().Truncate(24 * time.Hour).Format("1/02/2006") cTime = strings.ReplaceAll(cTime, today, "Today") aTime = strings.ReplaceAll(aTime, today, "Today") mTime = strings.ReplaceAll(mTime, today, "Today") @@ -522,6 +575,26 @@ func configMain(ctx context.Context, gptCliCtx *GptCliContext, args []string) er if err != nil { return fmt.Errorf("Could not write OpenAI API key file %v: %w", keyPath, err) } + sessPath := path.Join(configDir, SessionFile) + _, err = os.Stat(sessPath) + if err != nil && !os.IsNotExist(err) { + return fmt.Errorf("Could not open OpenAI Session file %v: %w", keyPath, err) + } + fmt.Printf("Enter your OpenAI Session key (optional): ") + sess, err := gptCliCtx.input.ReadString('\n') + if err != nil { + return err + } + sess = strings.TrimSpace(sess) + if len(sess) != 0 { + err = ioutil.WriteFile(sessPath, []byte(sess), 0600) + if err != nil { + return fmt.Errorf("Could not write OpenAI Session file %v: %w", keyPath, err) + } + gptCliCtx.haveSess = true + } else { + gptCliCtx.haveSess = false + } threadsPath := path.Join(configDir, ThreadsDir) err = os.MkdirAll(threadsPath, 0700) if err != nil { @@ -530,6 +603,9 @@ func configMain(ctx context.Context, gptCliCtx *GptCliContext, args []string) er } gptCliCtx.client = openai.NewClient(key) + if gptCliCtx.haveSess { + gptCliCtx.sessClient = openai.NewClient(sess) + } gptCliCtx.needConfig = false return nil @@ -552,6 +628,14 @@ func getKeyPath() (string, error) { return filepath.Join(configDir, KeyFile), nil } +func getSessPath() (string, error) { + configDir, err := getConfigDir() + if err != nil { + return "", err + } + return filepath.Join(configDir, SessionFile), nil +} + func getThreadsDir() (string, error) { configDir, err := getConfigDir() if err != nil { @@ -576,6 +660,22 @@ func loadKey() (string, error) { return string(data), nil } +func loadSess() (string, error) { + sessPath, err := getSessPath() + if err != nil { + return "", fmt.Errorf("Could not load OpenAI Session: %w", err) + } + data, err := ioutil.ReadFile(sessPath) + if err != nil { + if os.IsNotExist(err) { + return "", fmt.Errorf("Could not load OpenAI Session: "+ + "run `%v config` to configure", CommandName) + } + return "", fmt.Errorf("Could not load OpenAI Session: %w", err) + } + return string(data), nil +} + func getMultiLineInputRemainder(gptCliCtx *GptCliContext) (string, error) { var ret string var tmp string @@ -702,6 +802,8 @@ func splitBlocks(text string) []string { } func main() { + checkAndPrintUpgradeWarning() + ctx := context.Background() gptCliCtx := NewGptCliContext() diff --git a/cmd/gptcli/main_test.go b/cmd/gptcli/main_test.go index b8535cc..b4a3bc9 100644 --- a/cmd/gptcli/main_test.go +++ b/cmd/gptcli/main_test.go @@ -210,3 +210,70 @@ func TestGetCmdOrPrompt(t *testing.T) { }) } } + +func TestThreadSwitchMain(t *testing.T) { + ctrl := gomock.NewController(t) + defer ctrl.Finish() + + mockOpenAIClient := internal.NewMockOpenAIClient(ctrl) + + tmpFile, err := os.CreateTemp("", "gptcli.testThreadSwitch.*") + assert.Nil(t, err) + tmpFile.Close() + defer os.Remove(tmpFile.Name()) + + now := time.Now() + gptCliCtx := GptCliContext{ + client: mockOpenAIClient, + totThreads: 1, + needConfig: false, + curThreadNum: 1, + threads: []*GptCliThread{ + { + Dialogue: nil, + filePath: tmpFile.Name(), + ModTime: now, + AccessTime: now, + }, + }, + } + + tests := []struct { + name string + args []string + wantErr bool + errMsg string + newThread int + }{ + { + name: "successful thread switch", + args: []string{"thread", "1"}, + wantErr: false, + errMsg: "", + newThread: 1, + }, + { + name: "non-existent thread switch", + args: []string{"thread", "2"}, + wantErr: true, + errMsg: "Thread 2 does not exist. To list threads try 'ls'.\n", + newThread: 1, // No change expected + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + err := threadSwitchMain(context.Background(), &gptCliCtx, tt.args) + + if tt.wantErr { + assert.Error(t, err) + assert.Contains(t, err.Error(), tt.errMsg) + } else { + assert.NoError(t, err) + } + + // Verify the current thread number has been set correctly + assert.Equal(t, tt.newThread, gptCliCtx.curThreadNum) + }) + } +} diff --git a/go.mod b/go.mod index 887368d..888ca14 100644 --- a/go.mod +++ b/go.mod @@ -2,10 +2,12 @@ module github.com/mikeb26/gptcli go 1.21.4 +replace github.com/sashabaranov/go-openai => github.com/mikeb26/sashabaranov-go-openai v1.17.7-0.20231114224213-3c451cc966bc + require ( github.com/fatih/color v1.16.0 github.com/golang/mock v1.6.0 - github.com/sashabaranov/go-openai v1.17.5 + github.com/sashabaranov/go-openai v0.0.0-00010101000000-000000000000 github.com/stretchr/testify v1.8.4 ) @@ -14,9 +16,6 @@ require ( github.com/mattn/go-colorable v0.1.13 // indirect github.com/mattn/go-isatty v0.0.20 // indirect github.com/pmezard/go-difflib v1.0.0 // indirect - golang.org/x/mod v0.4.2 // indirect golang.org/x/sys v0.14.0 // indirect - golang.org/x/tools v0.1.1 // indirect - golang.org/x/xerrors v0.0.0-20200804184101-5ec99f83aff1 // indirect gopkg.in/yaml.v3 v3.0.1 // indirect ) diff --git a/go.sum b/go.sum index ff23d0f..bb09de5 100644 --- a/go.sum +++ b/go.sum @@ -9,16 +9,15 @@ github.com/mattn/go-colorable v0.1.13/go.mod h1:7S9/ev0klgBDR4GtXTXX8a3vIGJpMovk github.com/mattn/go-isatty v0.0.16/go.mod h1:kYGgaQfpe5nmfYZH+SKPsOc2e4SrIfOl2e/yFXSvRLM= github.com/mattn/go-isatty v0.0.20 h1:xfD0iDuEKnDkl03q4limB+vH+GxLEtL/jb4xVJSWWEY= github.com/mattn/go-isatty v0.0.20/go.mod h1:W+V8PltTTMOvKvAeJH7IuucS94S2C6jfK/D7dTCTo3Y= +github.com/mikeb26/sashabaranov-go-openai v1.17.7-0.20231114224213-3c451cc966bc h1:m7zAj1AlTPnTIzigHPN0bxnVybaR/iUnSH7uJhQTrTs= +github.com/mikeb26/sashabaranov-go-openai v1.17.7-0.20231114224213-3c451cc966bc/go.mod h1:lj5b/K+zjTSFxVLijLSTDZuP7adOgerWeFyZLUhAKRg= github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM= github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4= -github.com/sashabaranov/go-openai v1.17.5 h1:ItBzlrrfTtkFWOFlgfOhk3y/xRBC4PJol4gdbiK7hgg= -github.com/sashabaranov/go-openai v1.17.5/go.mod h1:lj5b/K+zjTSFxVLijLSTDZuP7adOgerWeFyZLUhAKRg= github.com/stretchr/testify v1.8.4 h1:CcVxjf3Q8PM0mHUKJCdn+eZZtm5yQwehR5yeSVQQcUk= github.com/stretchr/testify v1.8.4/go.mod h1:sz/lmYIOXD/1dqDmKjjqLyZ2RngseejIcXlSw2iwfAo= github.com/yuin/goldmark v1.3.5/go.mod h1:mwnBkeHKe2W/ZEtQ+71ViKU8L12m81fl3OWwC1Zlc8k= golang.org/x/crypto v0.0.0-20190308221718-c2843e01d9a2/go.mod h1:djNgcEr1/C05ACkg1iLfiJU5Ep61QUkGW8qpdssI0+w= golang.org/x/crypto v0.0.0-20191011191535-87dc89f01550/go.mod h1:yigFU9vqHzYiE8UmvKecakEJjdnWj3jj499lnFckfCI= -golang.org/x/mod v0.4.2 h1:Gz96sIWK3OalVv/I/qNygP42zyoKp3xptRVCWRFEBvo= golang.org/x/mod v0.4.2/go.mod h1:s0Qsj1ACt9ePp/hMypM3fl4fZqREWJwdYDEqhRiZZUA= golang.org/x/net v0.0.0-20190404232315-eb5bcb51f2a3/go.mod h1:t9HGtf8HONx5eT2rtn7q6eTqICYqUVnKs3thJo3Qplg= golang.org/x/net v0.0.0-20190620200207-3b0461eec859/go.mod h1:z5CRVTTTmAJ677TzLLGU+0bjPO0LkuOLi4/5GtJWs/s= @@ -39,11 +38,9 @@ golang.org/x/text v0.3.0/go.mod h1:NqM8EUOU14njkJ3fqMW+pc6Ldnwhi/IjpwHt7yyuwOQ= golang.org/x/text v0.3.3/go.mod h1:5Zoc/QRtKVWzQhOtBMvqHzDpF6irO9z98xDceosuGiQ= golang.org/x/tools v0.0.0-20180917221912-90fa682c2a6e/go.mod h1:n7NCudcB/nEzxVGmLbDWY5pfWTLqBcC2KZ6jyYvM4mQ= golang.org/x/tools v0.0.0-20191119224855-298f0cb1881e/go.mod h1:b+2E5dAYhXwXZwtnZ6UAqBI28+e2cm9otk0dWdXHAEo= -golang.org/x/tools v0.1.1 h1:wGiQel/hW0NnEkJUk8lbzkX2gFJU6PFxf1v5OlCfuOs= golang.org/x/tools v0.1.1/go.mod h1:o0xws9oXOQQZyjljx8fwUC0k7L1pTE6eaCbjGeHmOkk= golang.org/x/xerrors v0.0.0-20190717185122-a985d3407aa7/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0= golang.org/x/xerrors v0.0.0-20191011141410-1b5146add898/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0= -golang.org/x/xerrors v0.0.0-20200804184101-5ec99f83aff1 h1:go1bK/D/BFZV2I8cIQd1NKEZ+0owSTG1fDTci4IqFcE= golang.org/x/xerrors v0.0.0-20200804184101-5ec99f83aff1/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0= gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405 h1:yhCVgyC4o1eVCa2tZl7eS0r+SDo693bJlVdllGtEeKM= gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0= diff --git a/internal/interfaces.go b/internal/interfaces.go index 216729e..7eda1c5 100644 --- a/internal/interfaces.go +++ b/internal/interfaces.go @@ -6,6 +6,7 @@ package internal import ( "context" + "time" "github.com/sashabaranov/go-openai" ) @@ -13,4 +14,5 @@ import ( //go:generate mockgen --build_flags=--mod=mod -destination=openai_client_mock.go -package=$GOPACKAGE github.com/mikeb26/gptcli/internal OpenAIClient type OpenAIClient interface { CreateChatCompletion(ctx context.Context, request openai.ChatCompletionRequest) (response openai.ChatCompletionResponse, err error) + GetBillingUsage(ctx context.Context, startDate time.Time, endDate time.Time) (response openai.BillingUsageResponse, err error) }