fix: 修复余额查询的接口 (#185)

docs/userGuide.md

@@ -21,7 +21,7 @@
 |  **单聊**  | 每次对话都是一次新的对话，没有聊天上下文联系 |    <details><br /><summary>点击查看</summary><br /><img src="https://cdn.staticaly.com/gh/eryajf/tu/main/img/image_20230404_193608.jpg"><br /></details>                                                          |      |
 |  **串聊**  |            带上下文联系的对话模式            |     <details><br /><summary>点击查看</summary><br /><img src="https://cdn.staticaly.com/gh/eryajf/tu/main/img/image_20230404_193608.jpg"><br /></details>                                                         |      |
 |  **重置**  |        重置上下文模式，回归到默认模式        |        <details><br /><summary>点击查看</summary><br /><img src="https://cdn.staticaly.com/gh/eryajf/tu/main/img/image_20230404_193608.jpg"><br /></details>                                                      |      |
-|  **余额**  | ~~查询机器人所用OpenAI账号的余额~~(暂不可用) |       <details><br /><summary>点击查看</summary><br /><img src="https://cdn.staticaly.com/gh/eryajf/tu/main/img/image_20230304_222522.jpg"><br /></details>                                                       |      |
+|  **余额**  | 查询机器人所用OpenAI账号的余额 |       <details><br /><summary>点击查看</summary><br /><img src="https://cdn.staticaly.com/gh/eryajf/tu/main/img/image_20230304_222522.jpg"><br /></details>                                                       |      |
 |  **模板**  |           查看应用内置的prompt模板           |      <details><br /><summary>点击查看</summary><br /><img src="https://cdn.staticaly.com/gh/eryajf/tu/main/img/image_20230404_193827.jpg"><br /></details>                                                        |      |
 |  **图片**  |          查看如何根据提示生成图片          | <details><br /><summary>点击查看</summary><br /><img src="https://cdn.staticaly.com/gh/eryajf/tu/main/img/image_20230404_194125.jpg"><br /></details> |      |
 | **查对话** |            获取指定人员的对话历史            |      <details><br /><summary>点击查看</summary><br /><img src="https://cdn.staticaly.com/gh/eryajf/tu/main/img/image_20230404_193938.jpg"><br /></details>                                                        |      |

pkg/process/process_request.go

@@ -62,18 +62,16 @@ func ProcessRequest(rmsg *dingbot.ReceiveMsg) error {
 				logger.Warning(fmt.Errorf("send message error: %v", err))
 			}
 		case "余额":
-			// 	cacheMsg := public.UserService.GetUserMode("system_balance")
-			// 	if cacheMsg == "" {
-			// 		rst, err := public.GetBalance()
-			// 		if err != nil {
-			// 			logger.Warning(fmt.Errorf("get balance error: %v", err))
-			// 			return err
-			// 		}
-			// 		t1 := time.Unix(int64(rst.Grants.Data[0].EffectiveAt), 0)
-			// 		t2 := time.Unix(int64(rst.Grants.Data[0].ExpiresAt), 0)
-			// 		cacheMsg = fmt.Sprintf("💵 已用: 💲%v\n💵 剩余: 💲%v\n⏳ 有效时间: 从 %v 到 %v\n", fmt.Sprintf("%.2f", rst.TotalUsed), fmt.Sprintf("%.2f", rst.TotalAvailable), t1.Format("2006-01-02 15:04:05"), t2.Format("2006-01-02 15:04:05"))
-			// 	}
-			cacheMsg := "官方暂时改写了余额接口，因此暂不提供查询余额功能！2023-04-03"
+			cacheMsg := public.UserService.GetUserMode("system_balance")
+			if cacheMsg == "" {
+				rst, err := public.GetBalance()
+				if err != nil {
+					logger.Warning(fmt.Errorf("get balance error: %v", err))
+					return err
+				}
+				cacheMsg = rst
+			}
+			// cacheMsg := "官方暂时改写了余额接口，因此暂不提供查询余额功能！2023-04-03"
 			_, err := rmsg.ReplyToDingtalk(string(dingbot.TEXT), cacheMsg)
 			if err != nil {
 				logger.Warning(fmt.Errorf("send message error: %v", err))

public/balance.go

@@ -15,30 +15,88 @@ func InitAiCli() *resty.Client {
 	return resty.New().SetTimeout(10*time.Second).SetHeader("Authorization", fmt.Sprintf("Bearer %s", Config.ApiKey)).SetRetryCount(3).SetRetryWaitTime(2 * time.Second)
 }
 
-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 Bill struct {
+	Object     string      `json:"object"`
+	DailyCosts []DailyCost `json:"daily_costs"`
+	TotalUsage float64     `json:"total_usage"`
+}
+
+type DailyCost struct {
+	Timestamp float64    `json:"timestamp"`
+	LineItems []LineItem `json:"line_items"`
+}
+
+type LineItem struct {
+	Name string  `json:"name"`
+	Cost float64 `json:"cost"`
 }
 
 // GetBalance 获取账号余额
-func GetBalance() (Billing, error) {
-	var data Billing
-	var url string = "https://api.openai.com/dashboard/billing/credit_grants"
+func GetBalance() (string, error) {
+	var data Bill
+	path := "/v1/dashboard/billing/usage"
+	var url string = "https://api.openai.com" + path
 	if Config.BaseURL != "" {
-		url = Config.BaseURL + "/dashboard/billing/credit_grants"
+		url = Config.BaseURL + path
+	}
+	d, _ := time.ParseDuration("-24h")
+	resp, err := InitAiCli().R().SetQueryParams(map[string]string{
+		"start_date": time.Now().Add(d * 90).Format("2006-01-02"),
+		"end_date":   time.Now().Format("2006-01-02"),
+	}).Get(url)
+	if err != nil {
+		return "", err
+	}
+	err = json.Unmarshal(resp.Body(), &data)
+	if err != nil {
+		return "", err
+	}
+	sub, err := GetSub()
+	if err != nil {
+		return "", err
+	}
+	expireDate := time.Unix(sub.AccessUntil, 0).Format("2006-01-02 15:04:05")
+	used := data.TotalUsage / 100
+	totalAvailable := sub.HardLimitUsd - used
+	msg := fmt.Sprintf("💵 已用: 💲%v\n💵 剩余: 💲%v\n🕰 到期时间: %v\n", fmt.Sprintf("%.2f", used), fmt.Sprintf("%.2f", totalAvailable), expireDate)
+	// 放入缓存
+	UserService.SetUserMode("system_balance", msg)
+	return msg, nil
+}
+
+type Subscription struct {
+	Object             string      `json:"object"`
+	HasPaymentMethod   bool        `json:"has_payment_method"`
+	Canceled           bool        `json:"canceled"`
+	CanceledAt         interface{} `json:"canceled_at"`
+	Delinquent         interface{} `json:"delinquent"`
+	AccessUntil        int64       `json:"access_until"`
+	SoftLimit          int64       `json:"soft_limit"`
+	HardLimit          int64       `json:"hard_limit"`
+	SystemHardLimit    int64       `json:"system_hard_limit"`
+	SoftLimitUsd       float64     `json:"soft_limit_usd"`
+	HardLimitUsd       float64     `json:"hard_limit_usd"`
+	SystemHardLimitUsd float64     `json:"system_hard_limit_usd"`
+	Plan               Plan        `json:"plan"`
+	AccountName        string      `json:"account_name"`
+	PoNumber           interface{} `json:"po_number"`
+	BillingEmail       interface{} `json:"billing_email"`
+	TaxIDS             interface{} `json:"tax_ids"`
+	BillingAddress     interface{} `json:"billing_address"`
+	BusinessAddress    interface{} `json:"business_address"`
+}
+
+type Plan struct {
+	Title string `json:"title"`
+	ID    string `json:"id"`
+}
+
+func GetSub() (Subscription, error) {
+	var data Subscription
+	path := "/v1/dashboard/billing/subscription"
+	var url string = "https://api.openai.com" + path
+	if Config.BaseURL != "" {
+		url = Config.BaseURL + path
 	}
 	resp, err := InitAiCli().R().Get(url)
 	if err != nil {
@@ -48,10 +106,5 @@ func GetBalance() (Billing, error) {
 	if err != nil {
 		return data, err
 	}
-	t1 := time.Unix(int64(data.Grants.Data[0].EffectiveAt), 0)
-	t2 := time.Unix(int64(data.Grants.Data[0].ExpiresAt), 0)
-	msg := fmt.Sprintf("💵 已用: 💲%v\n💵 剩余: 💲%v\n⏳ 有效时间: 从 %v 到 %v\n", fmt.Sprintf("%.2f", data.TotalUsed), fmt.Sprintf("%.2f", data.TotalAvailable), t1.Format("2006-01-02 15:04:05"), t2.Format("2006-01-02 15:04:05"))
-	// 放入缓存
-	UserService.SetUserMode("system_balance", msg)
 	return data, nil
 }

public/public.go

@@ -4,6 +4,7 @@ import (
 	"github.com/eryajf/chatgpt-dingtalk/config"
 	"github.com/eryajf/chatgpt-dingtalk/pkg/cache"
 	"github.com/eryajf/chatgpt-dingtalk/pkg/db"
+	"github.com/sashabaranov/go-openai"
 )
 
 var UserService cache.UserServiceInterface
@@ -20,7 +21,7 @@ func InitSvc() {
 	// 初始化数据库
 	db.InitDB()
 	// 暂时不在初始化时获取余额
-	// if Config.Model == openai.GPT3Dot5Turbo0301 || Config.Model == openai.GPT3Dot5Turbo {
-	// _, _ = GetBalance()
-	// }
-}
+	if Config.Model == openai.GPT3Dot5Turbo0301 || Config.Model == openai.GPT3Dot5Turbo {
+		_, _ = GetBalance()
+	}
+}