feat: add balance billing query

Jimi6 · Mar 27, 2023 · 14c3aa4 · 14c3aa4
1 parent 5751762
commit 14c3aa4
Show file tree

Hide file tree

Showing 8 changed files with 130 additions and 6 deletions.
diff --git a/code/go.mod b/code/go.mod
@@ -7,6 +7,7 @@ require github.com/larksuite/oapi-sdk-go/v3 v3.0.14
 require (
 	github.com/gin-gonic/gin v1.8.2
 	github.com/google/uuid v1.3.0
+	github.com/k0kubun/pp/v3 v3.2.0
 	github.com/larksuite/oapi-sdk-gin v1.0.0
 	github.com/patrickmn/go-cache v2.1.0+incompatible
 	github.com/pion/opus v0.0.0-20230123082803-1052c3e89e58
@@ -25,6 +26,7 @@ require (
 	github.com/json-iterator/go v1.1.12 // indirect
 	github.com/leodido/go-urn v1.2.1 // indirect
 	github.com/magiconair/properties v1.8.7 // indirect
+	github.com/mattn/go-colorable v0.1.13 // indirect
 	github.com/mattn/go-isatty v0.0.17 // indirect
 	github.com/mitchellh/mapstructure v1.5.0 // indirect
 	github.com/modern-go/concurrent v0.0.0-20180306012644-bacd9c7ef1dd // indirect

diff --git a/code/go.sum b/code/go.sum
@@ -147,6 +147,8 @@ github.com/json-iterator/go v1.1.12 h1:PV8peI4a0ysnczrg+LtxykD8LfKY9ML6u2jnxaEnr
 github.com/json-iterator/go v1.1.12/go.mod h1:e30LSqwooZae/UwlEbR2852Gd8hjQvJoHmT4TnhNGBo=
 github.com/jstemmer/go-junit-report v0.0.0-20190106144839-af01ea7f8024/go.mod h1:6v2b51hI/fHJwM22ozAgKL4VKDeJcHhJFhtBdhmNjmU=
 github.com/jstemmer/go-junit-report v0.9.1/go.mod h1:Brl9GWCQeLvo8nXZwPNNblvFj/XSXhF0NWZEnDohbsk=
+github.com/k0kubun/pp/v3 v3.2.0 h1:h33hNTZ9nVFNP3u2Fsgz8JXiF5JINoZfFq4SvKJwNcs=
+github.com/k0kubun/pp/v3 v3.2.0/go.mod h1:ODtJQbQcIRfAD3N+theGCV1m/CBxweERz2dapdz1EwA=
 github.com/kisielk/gotool v1.0.0/go.mod h1:XhKaO+MFFWcvkIS/tQcRk01m1F5IRFswLeQ+oQHNcck=
 github.com/kr/fs v0.1.0/go.mod h1:FFnZGqtBN9Gxj7eW1uZ42v5BccTP0vu6NEaFoC2HwRg=
 github.com/kr/pretty v0.1.0/go.mod h1:dAy3ld7l9f0ibDNOQOHHMYYIIbhfbHSm3C4ZsoJORNo=
@@ -165,6 +167,9 @@ github.com/leodido/go-urn v1.2.1 h1:BqpAaACuzVSgi/VLzGZIobT2z4v53pjosyNd9Yv6n/w=
 github.com/leodido/go-urn v1.2.1/go.mod h1:zt4jvISO2HfUBqxjfIshjdMTYS56ZS/qv49ictyFfxY=
 github.com/magiconair/properties v1.8.7 h1:IeQXZAiQcpL9mgcAe1Nu6cX9LLw6ExEHKjN0VQdvPDY=
 github.com/magiconair/properties v1.8.7/go.mod h1:Dhd985XPs7jluiymwWYZ0G4Z61jb3vdS329zhj2hYo0=
+github.com/mattn/go-colorable v0.1.13 h1:fFA4WZxdEF4tXPZVKMLwD8oUnCTTo08duU7wxecdEvA=
+github.com/mattn/go-colorable v0.1.13/go.mod h1:7S9/ev0klgBDR4GtXTXX8a3vIGJpMovkB8vQcUbaXHg=
+github.com/mattn/go-isatty v0.0.16/go.mod h1:kYGgaQfpe5nmfYZH+SKPsOc2e4SrIfOl2e/yFXSvRLM=
 github.com/mattn/go-isatty v0.0.17 h1:BTarxUcIeDqL27Mc+vyvdWYSL28zpIhv3RoTdsLMPng=
 github.com/mattn/go-isatty v0.0.17/go.mod h1:kYGgaQfpe5nmfYZH+SKPsOc2e4SrIfOl2e/yFXSvRLM=
 github.com/mitchellh/mapstructure v1.5.0 h1:jeMsZIYE/09sWLaz43PL7Gy6RuMjD2eJVyuac5Z2hdY=

diff --git a/code/handlers/event_common_action.go b/code/handlers/event_common_action.go
@@ -113,3 +113,20 @@ func (*HelpAction) Execute(a *ActionInfo) bool {
 	}
 	return true
 }
+
+type BalanceAction struct { /*余额*/
+}
+
+func (*BalanceAction) Execute(a *ActionInfo) bool {
+	if _, foundBalance := utils.EitherTrimEqual(a.info.qParsed,
+		"/balance", "余额"); foundBalance {
+		balanceResp, err := a.handler.gpt.GetBalance()
+		if err != nil {
+			replyMsg(*a.ctx, "查询余额失败，请稍后再试", a.info.msgId)
+			return false
+		}
+		sendBalanceCard(*a.ctx, a.info.sessionId, *balanceResp)
+		return false
+	}
+	return true
+}
diff --git a/code/handlers/handler.go b/code/handlers/handler.go
@@ -205,6 +205,7 @@ func (m MessageHandler) msgReceivedHandler(ctx context.Context, event *larkim.P2
 		&EmptyAction{},           //空消息处理
 		&ClearAction{},           //清除消息处理
 		&HelpAction{},            //帮助处理
+		&BalanceAction{},         //余额处理
 		&RolePlayAction{},        //角色扮演处理
 		&MessageAction{},         //消息处理
 

diff --git a/code/handlers/msg.go b/code/handlers/msg.go
@@ -5,12 +5,12 @@ import (
 	"context"
 	"encoding/base64"
 	"fmt"
-	"start-feishubot/initialization"
-	"start-feishubot/services"
-
 	"github.com/google/uuid"
 	larkcard "github.com/larksuite/oapi-sdk-go/v3/card"
 	larkim "github.com/larksuite/oapi-sdk-go/v3/service/im/v1"
+	"start-feishubot/initialization"
+	"start-feishubot/services"
+	"start-feishubot/services/openai"
 )
 
 type CardKind string
@@ -684,3 +684,27 @@ func sendVarImageCard(ctx context.Context, imageKey string,
 	)
 	return nil
 }
+
+//TotalGranted   float64   `json:"total_granted"`
+//TotalUsed      float64   `json:"total_used"`
+//TotalAvailable float64   `json:"total_available"`
+//EffectiveAt    time.Time `json:"effective_at"`
+//ExpiresAt      time.Time `json:"expires_at"`
+func sendBalanceCard(ctx context.Context, msgId *string,
+	balance openai.BalanceResponse) {
+	newCard, _ := newSendCard(
+		withHeader("🎰️ 余额查询", larkcard.TemplateBlue),
+		withMainMd(fmt.Sprintf("总额度: %.2f$", balance.TotalGranted)),
+		withMainMd(fmt.Sprintf("已用额度: %.2f$", balance.TotalUsed)),
+		withMainMd(fmt.Sprintf("可用额度: %.2f$",
+			balance.TotalAvailable)),
+		withNote(fmt.Sprintf("有效期: %s - %s",
+			balance.EffectiveAt.Format("2006-01-02 15:04:05"),
+			balance.ExpiresAt.Format("2006-01-02 15:04:05"))),
+	)
+	replyCard(
+		ctx,
+		msgId,
+		newCard,
+	)
+}
diff --git a/code/services/openai/billing.go b/code/services/openai/billing.go
@@ -0,0 +1,63 @@
+package openai
+
+import (
+	"fmt"
+	"net/http"
+	"time"
+)
+
+//https://api.openai.com/dashboard/billing/credit_grants
+type Billing struct {
+	Object         string  `json:"object"`
+	TotalGranted   float64 `json:"total_granted"`
+	TotalUsed      float64 `json:"total_used"`
+	TotalAvailable float64 `json:"total_available"`
+	Grants         struct {
+		Object string `json:"object"`
+		Data   []struct {
+			Object      string  `json:"object"`
+			ID          string  `json:"id"`
+			GrantAmount float64 `json:"grant_amount"`
+			UsedAmount  float64 `json:"used_amount"`
+			EffectiveAt float64 `json:"effective_at"`
+			ExpiresAt   float64 `json:"expires_at"`
+		} `json:"data"`
+	} `json:"grants"`
+}
+
+type BalanceResponse struct {
+	TotalGranted   float64   `json:"total_granted"`
+	TotalUsed      float64   `json:"total_used"`
+	TotalAvailable float64   `json:"total_available"`
+	EffectiveAt    time.Time `json:"effective_at"`
+	ExpiresAt      time.Time `json:"expires_at"`
+}
+
+func (gpt *ChatGPT) GetBalance() (*BalanceResponse, error) {
+	var data Billing
+	err := gpt.sendRequestWithBodyType(
+		gpt.ApiUrl+"/dashboard/billing/credit_grants",
+		http.MethodGet,
+		nilBody,
+		nil,
+		&data,
+	)
+	if err != nil {
+		return nil, fmt.Errorf("failed to get billing data: %v", err)
+	}
+
+	balance := &BalanceResponse{
+		TotalGranted:   data.TotalGranted,
+		TotalUsed:      data.TotalUsed,
+		TotalAvailable: data.TotalAvailable,
+		ExpiresAt:      time.Now(),
+		EffectiveAt:    time.Now(),
+	}
+
+	if len(data.Grants.Data) > 0 {
+		balance.EffectiveAt = time.Unix(int64(data.Grants.Data[0].EffectiveAt), 0)
+		balance.ExpiresAt = time.Unix(int64(data.Grants.Data[0].ExpiresAt), 0)
+	}
+
+	return balance, nil
+}
diff --git a/code/services/openai/common.go b/code/services/openai/common.go
@@ -27,6 +27,8 @@ const (
 	jsonBody requestBodyType = iota
 	formVoiceDataBody
 	formPictureDataBody
+
+	nilBody
 )
 
 func (gpt ChatGPT) doAPIRequestWithRetry(url, method string, bodyType requestBodyType,
@@ -35,16 +37,15 @@ func (gpt ChatGPT) doAPIRequestWithRetry(url, method string, bodyType requestBod
 	var requestBodyData []byte
 	var err error
 	var writer *multipart.Writer
+	api = gpt.Lb.GetAPI()
 
 	switch bodyType {
 	case jsonBody:
-		api = gpt.Lb.GetAPI()
 		requestBodyData, err = json.Marshal(requestBody)
 		if err != nil {
 			return err
 		}
 	case formVoiceDataBody:
-		api = gpt.Lb.GetAPI()
 		formBody := &bytes.Buffer{}
 		writer = multipart.NewWriter(formBody)
 		err = audioMultipartForm(requestBody.(AudioToTextRequestBody), writer)
@@ -58,7 +59,6 @@ func (gpt ChatGPT) doAPIRequestWithRetry(url, method string, bodyType requestBod
 		requestBodyData = formBody.Bytes()
 
 	case formPictureDataBody:
-		api = gpt.Lb.GetAPI()
 		formBody := &bytes.Buffer{}
 		writer = multipart.NewWriter(formBody)
 		err = pictureMultipartForm(requestBody.(ImageVariantRequestBody), writer)
@@ -70,6 +70,8 @@ func (gpt ChatGPT) doAPIRequestWithRetry(url, method string, bodyType requestBod
 			return err
 		}
 		requestBodyData = formBody.Bytes()
+	case nilBody:
+		requestBodyData = nil
 
 	default:
 		return errors.New("unknown request body type")

diff --git a/code/services/openai/gpt3_test.go b/code/services/openai/gpt3_test.go
@@ -107,3 +107,13 @@ func TestVariateOneImageWithJpg(t *testing.T) {
 		t.Errorf("TestVariateOneImage returned empty imageURL")
 	}
 }
+
+func TestChatGPT_GetBalance(t *testing.T) {
+	config := initialization.LoadConfig("../../config.yaml")
+	gpt := NewChatGPT(*config)
+	balance, err := gpt.GetBalance()
+	if err != nil {
+		t.Errorf("TestChatGPT_GetBalance failed with error: %v", err)
+	}
+	fmt.Println("balance: ", balance)
+}