Ecosystem/ai-gateway
5
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Add: support update AIProxy balance

Browse Source
This commit is contained in:
adevjoe 2023-06-16 21:20:00 +08:00
parent 57b213a035
commit 0fb110a479
No known key found for this signature in database
GPG Key ID: 320204213030865E
1 changed files with 47 additions and 5 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

52
controller/channel-billing.go
View File

@ -4,13 +4,14 @@ import (
"encoding/json"
"errors"
"fmt"
"github.com/gin-gonic/gin"
"io"
"net/http"
"one-api/common"
"one-api/model"
"strconv"
"time"
"github.com/gin-gonic/gin"
)
// https://github.com/songquanpeng/one-api/issues/79
@ -44,12 +45,31 @@ type OpenAISBUsageResponse struct {
} `json:"data"`
}
func GetResponseBody(method, url string, channel *model.Channel) ([]byte, error) {
type AIProxyUserOverviewResponse struct {
Success bool `json:"success"`
Message string `json:"message"`
ErrorCode int `json:"error_code"`
Data struct {
TotalPoints float64 `json:"totalPoints"`
} `json:"data"`
}
// GetAuthHeader get auth header
func GetAuthHeader(token string) http.Header {
h := http.Header{}
h.Add("Authorization", fmt.Sprintf("Bearer %s", token))
return h
}
func GetResponseBody(method, url string, channel *model.Channel, headers http.Header) ([]byte, error) {
client := &http.Client{}
req, err := http.NewRequest(method, url, nil)
if err != nil {
return nil, err
}
for k := range headers {
req.Header.Add(k, headers.Get(k))
}
auth := fmt.Sprintf("Bearer %s", channel.Key)
req.Header.Add("Authorization", auth)
res, err := client.Do(req)
@ -69,7 +89,7 @@ func GetResponseBody(method, url string, channel *model.Channel) ([]byte, error)
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)
body, err := GetResponseBody("GET", url, channel, GetAuthHeader(channel.Key))
if err != nil {
return 0, err
}
@ -89,6 +109,26 @@ func updateChannelOpenAISBBalance(channel *model.Channel) (float64, error) {
return balance, nil
}
func updateChannelAIProxyBalance(channel *model.Channel) (float64, error) {
url := "https://aiproxy.io/api/report/getUserOverview"
headers := http.Header{}
headers.Add("Api-Key", channel.Key)
body, err := GetResponseBody("GET", url, channel, headers)
if err != nil {
return 0, err
}
response := AIProxyUserOverviewResponse{}
err = json.Unmarshal(body, &response)
if err != nil {
return 0, err
}
if !response.Success {
return 0, fmt.Errorf("code: %d, message: %s", response.ErrorCode, response.Message)
}
channel.UpdateBalance(response.Data.TotalPoints)
return response.Data.TotalPoints, nil
}
func updateChannelBalance(channel *model.Channel) (float64, error) {
baseURL := common.ChannelBaseURLs[channel.Type]
switch channel.Type {
@ -102,12 +142,14 @@ func updateChannelBalance(channel *model.Channel) (float64, error) {
baseURL = channel.BaseURL
case common.ChannelTypeOpenAISB:
return updateChannelOpenAISBBalance(channel)
case common.ChannelTypeAIProxy:
return updateChannelAIProxyBalance(channel)
default:
return 0, errors.New("尚未实现")
}
url := fmt.Sprintf("%s/v1/dashboard/billing/subscription", baseURL)
body, err := GetResponseBody("GET", url, channel)
body, err := GetResponseBody("GET", url, channel, GetAuthHeader(channel.Key))
if err != nil {
return 0, err
}
@ -123,7 +165,7 @@ func updateChannelBalance(channel *model.Channel) (float64, error) {
startDate = now.AddDate(0, 0, -100).Format("2006-01-02")
}
url = fmt.Sprintf("%s/v1/dashboard/billing/usage?start_date=%s&end_date=%s", baseURL, startDate, endDate)
body, err = GetResponseBody("GET", url, channel)
body, err = GetResponseBody("GET", url, channel, GetAuthHeader(channel.Key))
if err != nil {
return 0, err
}