feat: support balance query for CloseAI (#240)

* 增加CloseAI余额查询功能

* 去除debug代码

* fix: bug fix

* docs: update README

---------

Co-authored-by: sudongyang <sudongyang@xiaomi.com>
Co-authored-by: JustSong <songquanpeng@foxmail.com>
This commit is contained in:
Cross Fire
2023-07-03 20:43:42 +08:00
committed by GitHub
parent 04f40def2f
commit d383302e8a
4 changed files with 46 additions and 14 deletions

View File

@@ -32,6 +32,13 @@ type OpenAIUsageDailyCost struct {
}
}
type OpenAICreditGrants struct {
Object string `json:"object"`
TotalGranted float64 `json:"total_granted"`
TotalUsed float64 `json:"total_used"`
TotalAvailable float64 `json:"total_available"`
}
type OpenAIUsageResponse struct {
Object string `json:"object"`
//DailyCosts []OpenAIUsageDailyCost `json:"daily_costs"`
@@ -100,6 +107,22 @@ func GetResponseBody(method, url string, channel *model.Channel, headers http.He
return body, nil
}
func updateChannelCloseAIBalance(channel *model.Channel) (float64, error) {
url := fmt.Sprintf("%s/dashboard/billing/credit_grants", channel.BaseURL)
body, err := GetResponseBody("GET", url, channel, GetAuthHeader(channel.Key))
if err != nil {
return 0, err
}
response := OpenAICreditGrants{}
err = json.Unmarshal(body, &response)
if err != nil {
return 0, err
}
channel.UpdateBalance(response.TotalAvailable)
return response.TotalAvailable, nil
}
func updateChannelOpenAISBBalance(channel *model.Channel) (float64, error) {
url := fmt.Sprintf("https://api.openai-sb.com/sb-api/user/status?api_key=%s", channel.Key)
body, err := GetResponseBody("GET", url, channel, GetAuthHeader(channel.Key))
@@ -175,6 +198,9 @@ func updateChannelAIGC2DBalance(channel *model.Channel) (float64, error) {
func updateChannelBalance(channel *model.Channel) (float64, error) {
baseURL := common.ChannelBaseURLs[channel.Type]
if channel.BaseURL == "" {
channel.BaseURL = baseURL
}
switch channel.Type {
case common.ChannelTypeOpenAI:
if channel.BaseURL != "" {
@@ -184,6 +210,8 @@ func updateChannelBalance(channel *model.Channel) (float64, error) {
return 0, errors.New("尚未实现")
case common.ChannelTypeCustom:
baseURL = channel.BaseURL
case common.ChannelTypeCloseAI:
return updateChannelCloseAIBalance(channel)
case common.ChannelTypeOpenAISB:
return updateChannelOpenAISBBalance(channel)
case common.ChannelTypeAIProxy: