feat: support channel type AIGC2D (songquanpeng#220)

* feat: add AIGC2D Channel * chore: remove console logging & update balance rendering --------- Co-authored-by: Alone88 <[email protected]> Co-authored-by: JustSong <[email protected]>
knowlimit · Jun 29, 2023 · b1b3651 · b1b3651
1 parent 8f6bd51
commit b1b3651
Show file tree

Hide file tree

Showing 4 changed files with 31 additions and 1 deletion.
diff --git a/common/constants.go b/common/constants.go
@@ -148,6 +148,7 @@ const (
 	ChannelTypeAIProxy   = 10
 	ChannelTypePaLM      = 11
 	ChannelTypeAPI2GPT   = 12
+	ChannelTypeAIGC2D    = 13
 )
 
 var ChannelBaseURLs = []string{
@@ -164,4 +165,5 @@ var ChannelBaseURLs = []string{
 	"https://api.aiproxy.io",       // 10
 	"",                             // 11
 	"https://api.api2gpt.com",      // 12
+	"https://api.aigc2d.com",       // 13
 }
diff --git a/controller/channel-billing.go b/controller/channel-billing.go
@@ -61,6 +61,14 @@ type API2GPTUsageResponse struct {
 	TotalRemaining float64 `json:"total_remaining"`
 }
 
+type APGC2DGPTUsageResponse struct {
+	//Grants         interface{} `json:"grants"`
+	Object         string  `json:"object"`
+	TotalAvailable float64 `json:"total_available"`
+	TotalGranted   float64 `json:"total_granted"`
+	TotalUsed      float64 `json:"total_used"`
+}
+
 // GetAuthHeader get auth header
 func GetAuthHeader(token string) http.Header {
 	h := http.Header{}
@@ -150,6 +158,21 @@ func updateChannelAPI2GPTBalance(channel *model.Channel) (float64, error) {
 	return response.TotalRemaining, nil
 }
 
+func updateChannelAIGC2DBalance(channel *model.Channel) (float64, error) {
+	url := "https://api.aigc2d.com/dashboard/billing/credit_grants"
+	body, err := GetResponseBody("GET", url, channel, GetAuthHeader(channel.Key))
+	if err != nil {
+		return 0, err
+	}
+	response := APGC2DGPTUsageResponse{}
+	err = json.Unmarshal(body, &response)
+	if err != nil {
+		return 0, err
+	}
+	channel.UpdateBalance(response.TotalAvailable)
+	return response.TotalAvailable, nil
+}
+
 func updateChannelBalance(channel *model.Channel) (float64, error) {
 	baseURL := common.ChannelBaseURLs[channel.Type]
 	switch channel.Type {
@@ -167,6 +190,8 @@ func updateChannelBalance(channel *model.Channel) (float64, error) {
 		return updateChannelAIProxyBalance(channel)
 	case common.ChannelTypeAPI2GPT:
 		return updateChannelAPI2GPTBalance(channel)
+	case common.ChannelTypeAIGC2D:
+		return updateChannelAIGC2DBalance(channel)
 	default:
 		return 0, errors.New("尚未实现")
 	}

diff --git a/web/src/components/ChannelsTable.js b/web/src/components/ChannelsTable.js
@@ -38,6 +38,8 @@ function renderBalance(type, balance) {
       return <span>{renderNumber(balance)}</span>;
     case 12: // API2GPT
       return <span>¥{balance.toFixed(2)}</span>;
+    case 13: // AIGC2D
+      return <span>{renderNumber(balance)}</span>;
     default:
       return <span>不支持</span>;
   }

diff --git a/web/src/constants/channel.constants.js b/web/src/constants/channel.constants.js
@@ -9,5 +9,6 @@ export const CHANNEL_OPTIONS = [
   { key: 7, text: 'OhMyGPT', value: 7, color: 'purple' },
   { key: 9, text: 'AI.LS', value: 9, color: 'yellow' },
   { key: 10, text: 'AI Proxy', value: 10, color: 'purple' },
-  { key: 12, text: 'API2GPT', value: 12, color: 'blue' }
+  { key: 12, text: 'API2GPT', value: 12, color: 'blue' },
+  { key: 13, text: 'AIGC2D', value: 13, color: 'purple' }
 ];