Support Balance Query for DeepSeek

This commit is contained in:
Gk0Wk 2024-11-25 15:59:39 +08:00
parent 6ab87f8a08
commit 7d5c0959b2
3 changed files with 40 additions and 4 deletions

View File

@ -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"
)
@ -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"`
}
}
// GetAuthHeader get auth header
func GetAuthHeader(token string) http.Header {
h := http.Header{}
@ -245,6 +256,25 @@ 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
}
balance, err := strconv.ParseFloat(response.BalanceInfos[0].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() == "" {
@ -271,6 +301,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("尚未实现")
}

View File

@ -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:

View File

@ -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: