Jaymr/new-api
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

refactor: advanced custom channel route editor (#6865)

Browse Source 
* refactor: advanced custom channel route editor

* fix(channels): show raw balance response from balance cell
This commit is contained in:
Seefs
2026-08-18 17:31:21 +08:00
committed by GitHub
parent 3dda1d50c6
commit 2b0efd8484
21 changed files with 1551 additions and 506 deletions
Download Patch File Download Diff File Expand all files Collapse all files
controller/channel-billing.go
+130 -19
View File
@@ -5,13 +5,19 @@ import (
"errors"
"fmt"
"io"
"math"
"net/http"
"strconv"
"strings"
"time"
"github.com/QuantumNous/new-api/common"
"github.com/QuantumNous/new-api/constant"
"github.com/QuantumNous/new-api/model"
"github.com/QuantumNous/new-api/relay/channel/advancedcustom"
relaycommon "github.com/QuantumNous/new-api/relay/common"
relayconstant "github.com/QuantumNous/new-api/relay/constant"
"github.com/QuantumNous/new-api/relaykit/dto"
"github.com/QuantumNous/new-api/relaykit/types"
"github.com/QuantumNous/new-api/service"
"github.com/QuantumNous/new-api/setting/operation_setting"
@@ -47,6 +53,13 @@ type OpenAICreditGrants struct {
TotalAvailable float64 `json:"total_available"`
}
const maxAdvancedCustomBalanceResponseBytes = 256 << 10
type channelBalanceResult struct {
Balance float64
RawResponse string
}
type OpenAIUsageResponse struct {
Object string `json:"object"`
//DailyCosts []OpenAIUsageDailyCost `json:"daily_costs"`
@@ -174,7 +187,7 @@ func updateChannelCloseAIBalance(channel *model.Channel) (float64, error) {
return 0, err
}
response := OpenAICreditGrants{}
err = json.Unmarshal(body, &response)
err = common.Unmarshal(body, &response)
if err != nil {
return 0, err
}
@@ -189,7 +202,7 @@ func updateChannelOpenAISBBalance(channel *model.Channel) (float64, error) {
return 0, err
}
response := OpenAISBUsageResponse{}
err = json.Unmarshal(body, &response)
err = common.Unmarshal(body, &response)
if err != nil {
return 0, err
}
@@ -213,7 +226,7 @@ func updateChannelAIProxyBalance(channel *model.Channel) (float64, error) {
return 0, err
}
response := AIProxyUserOverviewResponse{}
err = json.Unmarshal(body, &response)
err = common.Unmarshal(body, &response)
if err != nil {
return 0, err
}
@@ -232,7 +245,7 @@ func updateChannelAPI2GPTBalance(channel *model.Channel) (float64, error) {
return 0, err
}
response := API2GPTUsageResponse{}
err = json.Unmarshal(body, &response)
err = common.Unmarshal(body, &response)
if err != nil {
return 0, err
}
@@ -247,7 +260,7 @@ func updateChannelSiliconFlowBalance(channel *model.Channel) (float64, error) {
return 0, err
}
response := SiliconFlowUsageResponse{}
err = json.Unmarshal(body, &response)
err = common.Unmarshal(body, &response)
if err != nil {
return 0, err
}
@@ -269,7 +282,7 @@ func updateChannelDeepSeekBalance(channel *model.Channel) (float64, error) {
return 0, err
}
response := DeepSeekUsageResponse{}
err = json.Unmarshal(body, &response)
err = common.Unmarshal(body, &response)
if err != nil {
return 0, err
}
@@ -298,7 +311,7 @@ func updateChannelAIGC2DBalance(channel *model.Channel) (float64, error) {
return 0, err
}
response := APGC2DGPTUsageResponse{}
err = json.Unmarshal(body, &response)
err = common.Unmarshal(body, &response)
if err != nil {
return 0, err
}
@@ -313,7 +326,7 @@ func updateChannelOpenRouterBalance(channel *model.Channel) (float64, error) {
return 0, err
}
response := OpenRouterCreditResponse{}
err = json.Unmarshal(body, &response)
err = common.Unmarshal(body, &response)
if err != nil {
return 0, err
}
@@ -343,7 +356,7 @@ func updateChannelMoonshotBalance(channel *model.Channel) (float64, error) {
}
response := MoonshotBalanceResponse{}
err = json.Unmarshal(body, &response)
err = common.Unmarshal(body, &response)
if err != nil {
return 0, err
}
@@ -356,7 +369,100 @@ func updateChannelMoonshotBalance(channel *model.Channel) (float64, error) {
return availableBalanceUsd, nil
}
func updateChannelBalance(channel *model.Channel) (float64, error) {
func fetchAdvancedCustomBalance(channel *model.Channel) (channelBalanceResult, error) {
key := strings.TrimSpace(channel.Key)
info := &relaycommon.RelayInfo{
RelayFormat: types.RelayFormatOpenAI,
RelayMode: relayconstant.RelayModeUnknown,
RequestURLPath: dto.AdvancedCustomBalancePath,
ChannelMeta: &relaycommon.ChannelMeta{
ChannelType: constant.ChannelTypeAdvancedCustom,
ChannelBaseUrl: channel.GetBaseURL(),
ApiKey: key,
ChannelOtherSettings: channel.GetOtherSettings(),
},
}
requestURL, headers, err := (&advancedcustom.Adaptor{}).BuildBalanceRequest(info)
if err != nil {
return channelBalanceResult{}, sanitizeFetchModelsError(err, key)
}
if err := applyFetchModelsHeaderOverrides(channel, key, headers); err != nil {
return channelBalanceResult{}, sanitizeFetchModelsError(err, key)
}
request, err := http.NewRequest(http.MethodGet, requestURL, nil)
if err != nil {
return channelBalanceResult{}, sanitizeFetchModelsError(err, key)
}
for name, values := range headers {
for _, value := range values {
request.Header.Add(name, value)
}
if strings.EqualFold(name, "Host") {
request.Host = headers.Get(name)
}
}
client, err := service.GetHttpClientWithProxy(channel.GetSetting().Proxy)
if err != nil {
return channelBalanceResult{}, sanitizeFetchModelsError(err, key)
}
response, err := client.Do(request)
if err != nil {
return channelBalanceResult{}, sanitizeAdvancedCustomRequestError(err, key, requestURL)
}
defer response.Body.Close()
if response.StatusCode != http.StatusOK {
return channelBalanceResult{}, fmt.Errorf("status code: %d", response.StatusCode)
}
body, err := io.ReadAll(io.LimitReader(response.Body, maxAdvancedCustomBalanceResponseBytes+1))
if err != nil {
return channelBalanceResult{}, sanitizeAdvancedCustomRequestError(err, key, requestURL)
}
if len(body) > maxAdvancedCustomBalanceResponseBytes {
return channelBalanceResult{}, fmt.Errorf("balance response exceeds %d bytes", maxAdvancedCustomBalanceResponseBytes)
}
var validated json.RawMessage
if err := common.Unmarshal(body, &validated); err != nil {
return channelBalanceResult{}, fmt.Errorf("invalid balance JSON response: %w", err)
}
if common.GetJsonType(validated) == "object" {
var creditSummary struct {
Object string `json:"object"`
TotalAvailable json.RawMessage `json:"total_available"`
}
if err := common.Unmarshal(body, &creditSummary); err != nil {
return channelBalanceResult{}, fmt.Errorf("invalid balance JSON response: %w", err)
}
if creditSummary.Object == "credit_summary" &&
common.GetJsonType(creditSummary.TotalAvailable) == "number" {
var balance float64
if err := common.Unmarshal(creditSummary.TotalAvailable, &balance); err == nil &&
balance >= 0 &&
!math.IsNaN(balance) &&
!math.IsInf(balance, 0) {
channel.UpdateBalance(balance)
return channelBalanceResult{Balance: balance}, nil
}
}
}
formatted, err := common.IndentJson(body)
if err != nil {
return channelBalanceResult{}, fmt.Errorf("invalid balance JSON response: %w", err)
}
return channelBalanceResult{RawResponse: string(formatted)}, nil
}
func updateChannelBalance(channel *model.Channel) (channelBalanceResult, error) {
if channel.Type == constant.ChannelTypeAdvancedCustom {
return fetchAdvancedCustomBalance(channel)
}
balance, err := updateStandardChannelBalance(channel)
return channelBalanceResult{Balance: balance}, err
}
func updateStandardChannelBalance(channel *model.Channel) (float64, error) {
baseURL := constant.ChannelBaseURLs[channel.Type]
if channel.GetBaseURL() == "" {
channel.BaseURL = &baseURL
@@ -396,7 +502,7 @@ func updateChannelBalance(channel *model.Channel) (float64, error) {
return 0, err
}
subscription := OpenAISubscriptionResponse{}
err = json.Unmarshal(body, &subscription)
err = common.Unmarshal(body, &subscription)
if err != nil {
return 0, err
}
@@ -412,7 +518,7 @@ func updateChannelBalance(channel *model.Channel) (float64, error) {
return 0, err
}
usage := OpenAIUsageResponse{}
err = json.Unmarshal(body, &usage)
err = common.Unmarshal(body, &usage)
if err != nil {
return 0, err
}
@@ -439,16 +545,21 @@ func UpdateChannelBalance(c *gin.Context) {
})
return
}
balance, err := updateChannelBalance(channel)
result, err := updateChannelBalance(channel)
if err != nil {
common.ApiError(c, err)
return
}
c.JSON(http.StatusOK, gin.H{
response := gin.H{
"success": true,
"message": "",
"balance": balance,
})
}
if result.RawResponse == "" {
response["balance"] = result.Balance
} else {
response["raw_response"] = result.RawResponse
}
c.JSON(http.StatusOK, response)
}
func updateAllChannelsBalance() error {
@@ -467,12 +578,12 @@ func updateAllChannelsBalance() error {
//if channel.Type != common.ChannelTypeOpenAI && channel.Type != common.ChannelTypeCustom {
// continue
//}
balance, err := updateChannelBalance(channel)
result, err := updateChannelBalance(channel)
if err != nil {
continue
} else {
} else if result.RawResponse == "" {
// err is nil & balance <= 0 means quota is used up
if balance <= 0 {
if result.Balance <= 0 {
service.DisableChannel(*types.NewChannelError(channel.Id, channel.Type, channel.Name, channel.ChannelInfo.IsMultiKey, "", channel.GetAutoBan()), "余额不足")
}
}
controller/channel_upstream_update.go
+29 -1
View File
@@ -304,6 +304,34 @@ func sanitizeFetchModelsError(err error, key string) error {
return errors.New(message)
}
func sanitizeAdvancedCustomRequestError(err error, key string, requestURL string) error {
err = sanitizeFetchModelsError(err, key)
if err == nil {
return nil
}
parsedURL, parseErr := url.Parse(requestURL)
if parseErr != nil {
return err
}
message := err.Error()
for _, value := range parsedURL.Query() {
for _, secret := range value {
if secret == "" {
continue
}
message = strings.ReplaceAll(message, secret, "[REDACTED]")
message = strings.ReplaceAll(message, url.QueryEscape(secret), "[REDACTED]")
message = strings.ReplaceAll(message, url.PathEscape(secret), "[REDACTED]")
}
}
if key != "" {
message = strings.ReplaceAll(message, key, "[REDACTED]")
message = strings.ReplaceAll(message, url.QueryEscape(key), "[REDACTED]")
message = strings.ReplaceAll(message, url.PathEscape(key), "[REDACTED]")
}
return errors.New(message)
}
func getFetchModelsResponseBody(method string, requestURL string, channel *model.Channel, headers http.Header) ([]byte, error) {
request, err := http.NewRequest(method, requestURL, nil)
if err != nil {
@@ -409,7 +437,7 @@ func fetchChannelUpstreamModelIDs(channel *model.Channel) ([]string, error) {
body, err := getFetchModelsResponseBody(http.MethodGet, url, channel, headers)
if err != nil {
return nil, sanitizeFetchModelsError(err, key)
return nil, sanitizeAdvancedCustomRequestError(err, key, url)
}
var result OpenAIModelsResponse
controller/channel_upstream_update_test.go
+9
View File
@@ -168,6 +168,15 @@ func TestFetchAdvancedCustomModelsRedactsQueryKeyFromTransportErrors(t *testing.
Err: errors.New("connection refused"),
}, secret)
require.EqualError(t, direct, "connection refused")
queryValue := "prefix-" + secret
queryError := sanitizeAdvancedCustomRequestError(
errors.New("dial "+queryValue+": connection refused"),
queryValue,
baseURL+"/v1/models?custom-token="+url.QueryEscape(queryValue),
)
require.NotContains(t, queryError.Error(), queryValue)
require.EqualError(t, queryError, "dial [REDACTED]: connection refused")
}
func TestFetchOrdinaryOpenAIModelsKeepsExistingEmptyDataBehavior(t *testing.T) {