Add support for querying billing/usage

.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

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
 

cmd/gptcli/help.txt

@@ -11,4 +11,4 @@ Available Commands:
   ls                             List available threads(conversations)
   thread <thread#>               Switch to a previously created thread
   exit                           Exit gptcli
-
+  billing                        Show billing/daily usage information

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()
 

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)
+		})
+	}
+}

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
 )

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=

internal/interfaces.go

@@ -6,11 +6,13 @@ package internal
 
 import (
 	"context"
+	"time"
 
 	"github.com/sashabaranov/go-openai"
 )
 
 //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)
 }