Skip to content

Commit

Permalink
feat: add balance query support for DeepSeek (#1946)
Browse files Browse the repository at this point in the history
* Support Balance Query for DeepSeek

* Fix
  • Loading branch information
Gk0Wk authored Dec 22, 2024
1 parent 381172c commit 95e8c16
Show file tree
Hide file tree
Showing 3 changed files with 50 additions and 4 deletions.
50 changes: 46 additions & 4 deletions controller/channel-billing.go
Original file line number Diff line number Diff line change
Expand Up @@ -4,16 +4,17 @@ import (
"encoding/json"
"errors"
"fmt"
"io"
"net/http"
"strconv"
"time"

"github.com/songquanpeng/one-api/common/client"
"github.com/songquanpeng/one-api/common/config"
"github.com/songquanpeng/one-api/common/logger"
"github.com/songquanpeng/one-api/model"
"github.com/songquanpeng/one-api/monitor"
"github.com/songquanpeng/one-api/relay/channeltype"
"io"
"net/http"
"strconv"
"time"

"github.com/gin-gonic/gin"
)
Expand Down Expand Up @@ -101,6 +102,16 @@ type SiliconFlowUsageResponse struct {
} `json:"data"`
}

type DeepSeekUsageResponse struct {
IsAvailable bool `json:"is_available"`
BalanceInfos []struct {
Currency string `json:"currency"`
TotalBalance string `json:"total_balance"`
GrantedBalance string `json:"granted_balance"`
ToppedUpBalance string `json:"topped_up_balance"`
} `json:"balance_infos"`
}

// GetAuthHeader get auth header
func GetAuthHeader(token string) http.Header {
h := http.Header{}
Expand Down Expand Up @@ -245,6 +256,35 @@ func updateChannelSiliconFlowBalance(channel *model.Channel) (float64, error) {
return balance, nil
}

func updateChannelDeepSeekBalance(channel *model.Channel) (float64, error) {
url := "https://api.deepseek.com/user/balance"
body, err := GetResponseBody("GET", url, channel, GetAuthHeader(channel.Key))
if err != nil {
return 0, err
}
response := DeepSeekUsageResponse{}
err = json.Unmarshal(body, &response)
if err != nil {
return 0, err
}
index := -1
for i, balanceInfo := range response.BalanceInfos {
if balanceInfo.Currency == "CNY" {
index = i
break
}
}
if index == -1 {
return 0, errors.New("currency CNY not found")
}
balance, err := strconv.ParseFloat(response.BalanceInfos[index].TotalBalance, 64)
if err != nil {
return 0, err
}
channel.UpdateBalance(balance)
return balance, nil
}

func updateChannelBalance(channel *model.Channel) (float64, error) {
baseURL := channeltype.ChannelBaseURLs[channel.Type]
if channel.GetBaseURL() == "" {
Expand All @@ -271,6 +311,8 @@ func updateChannelBalance(channel *model.Channel) (float64, error) {
return updateChannelAIGC2DBalance(channel)
case channeltype.SiliconFlow:
return updateChannelSiliconFlowBalance(channel)
case channeltype.DeepSeek:
return updateChannelDeepSeekBalance(channel)
default:
return 0, errors.New("尚未实现")
}
Expand Down
2 changes: 2 additions & 0 deletions web/berry/src/views/Channel/component/TableRow.js
Original file line number Diff line number Diff line change
Expand Up @@ -268,6 +268,8 @@ function renderBalance(type, balance) {
return <span>¥{balance.toFixed(2)}</span>;
case 13: // AIGC2D
return <span>{renderNumber(balance)}</span>;
case 36: // DeepSeek
return <span>¥{balance.toFixed(2)}</span>;
case 44: // SiliconFlow
return <span>¥{balance.toFixed(2)}</span>;
default:
Expand Down
2 changes: 2 additions & 0 deletions web/default/src/components/ChannelsTable.js
Original file line number Diff line number Diff line change
Expand Up @@ -52,6 +52,8 @@ function renderBalance(type, balance) {
return <span>¥{balance.toFixed(2)}</span>;
case 13: // AIGC2D
return <span>{renderNumber(balance)}</span>;
case 36: // DeepSeek
return <span>¥{balance.toFixed(2)}</span>;
case 44: // SiliconFlow
return <span>¥{balance.toFixed(2)}</span>;
default:
Expand Down

0 comments on commit 95e8c16

Please sign in to comment.