89 lines
3.5 KiB
Go
89 lines
3.5 KiB
Go
package relay
|
|
|
|
import (
|
|
"errors"
|
|
"fmt"
|
|
"net/http"
|
|
"strings"
|
|
|
|
"github.com/QuantumNous/new-api/common"
|
|
"github.com/QuantumNous/new-api/constant"
|
|
"github.com/QuantumNous/new-api/logger"
|
|
relaycommon "github.com/QuantumNous/new-api/relay/common"
|
|
"github.com/QuantumNous/new-api/relay/helper"
|
|
"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"
|
|
"github.com/gin-gonic/gin"
|
|
"github.com/samber/lo"
|
|
)
|
|
|
|
// PrepareRequestBilling estimates and reserves one request's charge. Transports
|
|
// provide the current request body through BodyStorage or BillingRequestInput;
|
|
// channel retries retain the resulting billing session and pricing snapshot.
|
|
func PrepareRequestBilling(c *gin.Context, info *relaycommon.RelayInfo) *types.NewAPIError {
|
|
durable := info.OriginModelName == "gpt-6-astra" && c.Request.Method == http.MethodPost && c.Request.URL.Path == "/v1/responses"
|
|
needSensitiveCheck := setting.ShouldCheckPromptSensitive()
|
|
meta := &types.TokenCountMeta{TokenType: types.TokenTypeTokenizer}
|
|
if info.Request != nil && (needSensitiveCheck || constant.CountToken || durable) {
|
|
meta = info.Request.GetTokenCountMeta()
|
|
} else {
|
|
// Avoid building CombineText when only the pricing quantities are needed.
|
|
switch request := info.Request.(type) {
|
|
case *dto.GeneralOpenAIRequest:
|
|
meta.MaxTokens = int(max(lo.FromPtr(request.MaxTokens), lo.FromPtr(request.MaxCompletionTokens)))
|
|
case *dto.OpenAIResponsesRequest:
|
|
meta.MaxTokens = int(lo.FromPtr(request.MaxOutputTokens))
|
|
case *dto.ClaudeRequest:
|
|
meta.MaxTokens = int(lo.FromPtr(request.MaxTokens))
|
|
case *dto.ImageRequest:
|
|
meta = request.GetTokenCountMeta()
|
|
}
|
|
}
|
|
|
|
if needSensitiveCheck && meta != nil {
|
|
if contains, words := service.CheckSensitiveText(meta.CombineText); contains {
|
|
service.RequestPolicy(c).AddEvent(service.PolicyEvent{ErrorCode: string(types.ErrorCodeSensitiveWordsDetected), ErrorSource: "local", Decision: service.PolicyDecision{Action: "stop", Reason: "local_rejection", Source: "global"}, Health: "unchanged"})
|
|
message := fmt.Sprintf("user sensitive words detected: %s", strings.Join(words, ", "))
|
|
logger.LogWarn(c, message)
|
|
return types.NewError(errors.New(message), types.ErrorCodeSensitiveWordsDetected)
|
|
}
|
|
}
|
|
|
|
tokens, err := service.EstimateRequestToken(c, meta, info)
|
|
if durable { tokens,err=service.CountRequestToken(c,meta,info) }
|
|
if err != nil {
|
|
return types.NewError(err, types.ErrorCodeCountTokenFailed)
|
|
}
|
|
info.SetEstimatePromptTokens(tokens)
|
|
|
|
priceData, err := helper.ModelPriceHelper(c, info, tokens, meta)
|
|
if err != nil {
|
|
return types.NewError(err, types.ErrorCodeModelPriceError, types.ErrOptionWithStatusCode(http.StatusBadRequest))
|
|
}
|
|
if durable {
|
|
info.DurableAttemptBilling = true
|
|
return nil
|
|
}
|
|
if priceData.FreeModel {
|
|
logger.LogInfo(c, common.LogText("model %s is free, skipping pre-consume", info.OriginModelName))
|
|
return nil
|
|
}
|
|
return service.PreConsumeBilling(c, priceData.QuotaToPreConsume, info)
|
|
}
|
|
|
|
// RefundFailedRequestBilling applies the common final-failure policy after all
|
|
// eligible attempts have ended. A settled BillingSession never refunds again.
|
|
func RefundFailedRequestBilling(c *gin.Context, info *relaycommon.RelayInfo, apiErr *types.NewAPIError) *types.NewAPIError {
|
|
if apiErr == nil {
|
|
return nil
|
|
}
|
|
apiErr = service.NormalizeViolationFeeError(apiErr)
|
|
if info.Billing != nil {
|
|
info.Billing.Refund(c)
|
|
}
|
|
service.ChargeViolationFeeIfNeeded(c, info, apiErr)
|
|
return apiErr
|
|
}
|