Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
13 changes: 13 additions & 0 deletions controller/model_list_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -267,10 +267,16 @@ func TestGetUserModelsExpandsAutoGroupsInConfiguredOrder(t *testing.T) {

func TestListModelsIncludesTieredBillingModel(t *testing.T) {
withSelfUseModeDisabled(t)
savedModelRatios := ratio_setting.ModelRatio2JSONString()
t.Cleanup(func() {
require.NoError(t, ratio_setting.UpdateModelRatioByJSONString(savedModelRatios))
})
require.NoError(t, ratio_setting.UpdateModelRatioByJSONString(`{"zz-utf8-visible-model":7.5}`))
withTieredBillingConfig(t, map[string]string{
"zz-tiered-visible-model": "tiered_expr",
"zz-tiered-empty-expr-model": "tiered_expr",
"zz-tiered-missing-expr-model": "tiered_expr",
"zz-utf8-visible-model": "utf8_bytes",
}, map[string]string{
"zz-tiered-visible-model": `tier("base", p * 1 + c * 2)`,
"zz-tiered-empty-expr-model": " ",
Expand All @@ -288,6 +294,7 @@ func TestListModelsIncludesTieredBillingModel(t *testing.T) {
{Group: "default", Model: "zz-tiered-visible-model", ChannelId: 1, Enabled: true},
{Group: "default", Model: "zz-tiered-empty-expr-model", ChannelId: 1, Enabled: true},
{Group: "default", Model: "zz-tiered-missing-expr-model", ChannelId: 1, Enabled: true},
{Group: "default", Model: "zz-utf8-visible-model", ChannelId: 1, Enabled: true},
{Group: "default", Model: "zz-unpriced-model", ChannelId: 1, Enabled: true},
}).Error)

Expand All @@ -302,6 +309,7 @@ func TestListModelsIncludesTieredBillingModel(t *testing.T) {
require.Contains(t, ids, "zz-tiered-visible-model")
require.NotContains(t, ids, "zz-tiered-empty-expr-model")
require.NotContains(t, ids, "zz-tiered-missing-expr-model")
require.Contains(t, ids, "zz-utf8-visible-model")
require.NotContains(t, ids, "zz-unpriced-model")

pricingByName := pricingByModelName(model.GetPricing())
Expand All @@ -319,6 +327,11 @@ func TestListModelsIncludesTieredBillingModel(t *testing.T) {
require.True(t, ok)
require.Empty(t, missingExprPricing.BillingMode)
require.Empty(t, missingExprPricing.BillingExpr)

utf8Pricing, ok := pricingByName["zz-utf8-visible-model"]
require.True(t, ok)
require.Equal(t, "utf8_bytes", utf8Pricing.BillingMode)
require.Empty(t, utf8Pricing.BillingExpr)
}

func TestListModelsUsesAdvancedCustomEndpointTypesFromPricingCache(t *testing.T) {
Expand Down
4 changes: 3 additions & 1 deletion controller/ratio_sync.go
Original file line number Diff line number Diff line change
Expand Up @@ -413,7 +413,9 @@ func FetchUpstreamRatios(c *gin.Context) {
if item.ModelName == "" {
continue
}
if item.BillingMode == billing_setting.BillingModeTieredExpr && strings.TrimSpace(item.BillingExpr) != "" {
if item.BillingMode == billing_setting.BillingModeUTF8Bytes {
billingModeMap[item.ModelName] = billing_setting.BillingModeUTF8Bytes
} else if item.BillingMode == billing_setting.BillingModeTieredExpr && strings.TrimSpace(item.BillingExpr) != "" {
billingModeMap[item.ModelName] = billing_setting.BillingModeTieredExpr
billingExprMap[item.ModelName] = item.BillingExpr
}
Expand Down
4 changes: 3 additions & 1 deletion controller/relay.go
Original file line number Diff line number Diff line change
Expand Up @@ -24,6 +24,7 @@ import (
"github.com/QuantumNous/new-api/relaykit/types"
"github.com/QuantumNous/new-api/service"
"github.com/QuantumNous/new-api/setting"
"github.com/QuantumNous/new-api/setting/billing_setting"
"github.com/QuantumNous/new-api/setting/operation_setting"

"github.com/bytedance/gopkg/util/gopool"
Expand Down Expand Up @@ -125,9 +126,10 @@ func Relay(c *gin.Context, relayFormat types.RelayFormat) {
newAPIError = types.NewError(err, types.ErrorCodeGenRelayInfoFailed)
return
}
relayInfo.BillingMode = billing_setting.GetBillingMode(relayInfo.OriginModelName)

needSensitiveCheck := setting.ShouldCheckPromptSensitive()
needCountToken := constant.CountToken
needCountToken := constant.CountToken || relayInfo.BillingMode == billing_setting.BillingModeUTF8Bytes
// Avoid building huge CombineText (strings.Join) when token counting and sensitive check are both disabled.
var meta *types.TokenCountMeta
if needSensitiveCheck || needCountToken {
Expand Down
4 changes: 3 additions & 1 deletion model/pricing.go
Original file line number Diff line number Diff line change
Expand Up @@ -400,7 +400,9 @@ func updatePricing() {
audioCompletionRatio := ratio_setting.GetAudioCompletionRatio(model)
pricing.AudioCompletionRatio = &audioCompletionRatio
}
if billingMode := billing_setting.GetBillingMode(model); billingMode == "tiered_expr" {
if billingMode := billing_setting.GetBillingMode(model); billingMode == billing_setting.BillingModeUTF8Bytes {
pricing.BillingMode = billingMode
} else if billingMode == billing_setting.BillingModeTieredExpr {
if expr, ok := billing_setting.GetBillingExpr(model); ok && strings.TrimSpace(expr) != "" {
pricing.BillingMode = billingMode
pricing.BillingExpr = expr
Expand Down
2 changes: 2 additions & 0 deletions relay/common/relay_info.go
Original file line number Diff line number Diff line change
Expand Up @@ -98,6 +98,7 @@ type RelayInfo struct {
UsePrice bool
RelayMode int
OriginModelName string
BillingMode string
RequestURLPath string
RequestHeaders map[string]string
ShouldIncludeUsage bool
Expand Down Expand Up @@ -263,6 +264,7 @@ func (info *RelayInfo) ToString() string {
fmt.Fprintf(b, "IsPlayground: %t, ", info.IsPlayground)
fmt.Fprintf(b, "RequestURLPath: %q, ", info.RequestURLPath)
fmt.Fprintf(b, "OriginModelName: %q, ", info.OriginModelName)
fmt.Fprintf(b, "BillingMode: %q, ", info.BillingMode)
fmt.Fprintf(b, "EstimatePromptTokens: %d, ", info.estimatePromptTokens)
fmt.Fprintf(b, "ShouldIncludeUsage: %t, ", info.ShouldIncludeUsage)
fmt.Fprintf(b, "DisablePing: %t, ", info.DisablePing)
Expand Down
14 changes: 11 additions & 3 deletions relay/helper/price.go
Original file line number Diff line number Diff line change
Expand Up @@ -72,11 +72,16 @@ func HandleGroupRatio(ctx *gin.Context, relayInfo *relaycommon.RelayInfo) hostty

func ModelPriceHelper(c *gin.Context, info *relaycommon.RelayInfo, promptTokens int, meta *types.TokenCountMeta) (hosttypes.PriceData, error) {
modelPrice, usePrice := ratio_setting.GetModelPrice(info.OriginModelName, false)
billingMode := info.BillingMode
if billingMode == "" {
billingMode = billing_setting.GetBillingMode(info.OriginModelName)
info.BillingMode = billingMode
}

groupRatioInfo := HandleGroupRatio(c, info)

// Check if this model uses tiered_expr billing
if billing_setting.GetBillingMode(info.OriginModelName) == billing_setting.BillingModeTieredExpr {
if billingMode == billing_setting.BillingModeTieredExpr {
return modelPriceHelperTiered(c, info, promptTokens, meta, groupRatioInfo)
}

Expand All @@ -92,8 +97,11 @@ func ModelPriceHelper(c *gin.Context, info *relaycommon.RelayInfo, promptTokens
var audioCompletionRatio float64
var freeModel bool
if !usePrice {
preConsumedTokens := common.Max(promptTokens, common.PreConsumedQuota)
if meta.MaxTokens != 0 {
preConsumedTokens := promptTokens
if billingMode != billing_setting.BillingModeUTF8Bytes {
preConsumedTokens = common.Max(promptTokens, common.PreConsumedQuota)
}
if billingMode != billing_setting.BillingModeUTF8Bytes && meta.MaxTokens != 0 {
preConsumedTokens += meta.MaxTokens
}
var success bool
Expand Down
38 changes: 38 additions & 0 deletions relay/helper/price_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -272,3 +272,41 @@ func TestModelPriceHelperRequestBillingRatiosOnlyApplyToFixedPrice(t *testing.T)
require.Equal(t, common.QuotaClampOverflow, clamp.Kind)
require.Nil(t, info.Billing)
}

func TestModelPriceHelperUTF8BytesDoesNotReserveCompletionTokens(t *testing.T) {
gin.SetMode(gin.TestMode)

saved := map[string]string{}
require.NoError(t, config.GlobalConfig.SaveToDB(func(key, value string) error {
saved[key] = value
return nil
}))
savedModelRatios := ratio_setting.ModelRatio2JSONString()
t.Cleanup(func() {
require.NoError(t, config.GlobalConfig.LoadFromDB(saved))
require.NoError(t, ratio_setting.UpdateModelRatioByJSONString(savedModelRatios))
})

require.NoError(t, config.GlobalConfig.LoadFromDB(map[string]string{
"billing_setting.billing_mode": `{"utf8-bytes-price-model":"utf8_bytes"}`,
"group_ratio_setting.group_ratio": `{"default":1}`,
}))
require.NoError(t, ratio_setting.UpdateModelRatioByJSONString(`{"utf8-bytes-price-model":7.5}`))

ctx, _ := gin.CreateTestContext(httptest.NewRecorder())
ctx.Set("group", "default")
info := &relaycommon.RelayInfo{
OriginModelName: "utf8-bytes-price-model",
BillingMode: billing_setting.BillingModeUTF8Bytes,
UserGroup: "default",
UsingGroup: "default",
}
require.NoError(t, config.GlobalConfig.LoadFromDB(map[string]string{
"billing_setting.billing_mode": `{}`,
}))

priceData, err := ModelPriceHelper(ctx, info, 7, &types.TokenCountMeta{MaxTokens: 1000})

require.NoError(t, err)
require.Equal(t, 52, priceData.QuotaToPreConsume)
}
28 changes: 26 additions & 2 deletions service/text_quota.go
Original file line number Diff line number Diff line change
Expand Up @@ -17,6 +17,7 @@ import (
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/setting/billing_setting"
"github.com/QuantumNous/new-api/setting/operation_setting"

"github.com/bytedance/gopkg/util/gopool"
Expand Down Expand Up @@ -263,8 +264,20 @@ func calculateTextQuotaSummary(ctx *gin.Context, relayInfo *relaycommon.RelayInf
summary.CacheCreationTokens1h = usage.ClaudeCacheCreation1hTokens
summary.ImageTokens = usage.PromptTokensDetails.ImageTokens
summary.AudioTokens = usage.PromptTokensDetails.AudioTokens
isUTF8BytesBilling := relayInfo.BillingMode == billing_setting.BillingModeUTF8Bytes
if isUTF8BytesBilling {
summary.PromptTokens = relayInfo.GetEstimatePromptTokens()
summary.CompletionTokens = 0
summary.TotalTokens = summary.PromptTokens
summary.CacheTokens = 0
summary.CacheCreationTokens = 0
summary.CacheCreationTokens5m = 0
summary.CacheCreationTokens1h = 0
summary.ImageTokens = 0
summary.AudioTokens = 0
}
legacyClaudeDerived := isLegacyClaudeDerivedOpenAIUsage(relayInfo, usage)
isOpenRouterClaudeBilling := relayInfo.ChannelMeta != nil &&
isOpenRouterClaudeBilling := !isUTF8BytesBilling && relayInfo.ChannelMeta != nil &&
relayInfo.ChannelType == constant.ChannelTypeOpenRouter &&
summary.IsClaudeUsageSemantic

Expand Down Expand Up @@ -476,7 +489,10 @@ func PostTextConsumeQuota(ctx *gin.Context, relayInfo *relaycommon.RelayInfo, us
} else {
other = GenerateTextOtherInfo(ctx, relayInfo, summary.ModelRatio, summary.GroupRatio, summary.CompletionRatio, summary.CacheTokens, summary.CacheRatio, summary.ModelPrice, relayInfo.PriceData.GroupRatioInfo.GroupSpecialRatio)
}
appendUsageBillingPathForLog(other, common.GetContextKeyBool(ctx, constant.ContextKeyLocalCountTokens), originUsage)
if relayInfo.BillingMode == billing_setting.BillingModeUTF8Bytes {
other["billing_mode"] = billing_setting.BillingModeUTF8Bytes
}
appendTextUsageBillingPathForLog(ctx, other, relayInfo, originUsage)
if adminRejectReason != "" {
other["reject_reason"] = adminRejectReason
}
Expand Down Expand Up @@ -541,3 +557,11 @@ func PostTextConsumeQuota(ctx *gin.Context, relayInfo *relaycommon.RelayInfo, us
perfmetrics.RecordRelaySample(relayInfo, true, int64(summary.CompletionTokens))
})
}

func appendTextUsageBillingPathForLog(ctx *gin.Context, other map[string]interface{}, relayInfo *relaycommon.RelayInfo, usage *dto.Usage) {
if relayInfo.BillingMode == billing_setting.BillingModeUTF8Bytes {
appendUsageBillingPathForLog(other, true, nil)
return
}
appendUsageBillingPathForLog(other, common.GetContextKeyBool(ctx, constant.ContextKeyLocalCountTokens), usage)
}
18 changes: 14 additions & 4 deletions service/token_counter.go
Original file line number Diff line number Diff line change
Expand Up @@ -15,6 +15,7 @@ import (
constant2 "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/setting/billing_setting"

"github.com/gin-gonic/gin"
)
Expand Down Expand Up @@ -177,15 +178,24 @@ func getImageToken(c *gin.Context, fileMeta *types.FileMeta, model string, strea
}

func EstimateRequestToken(c *gin.Context, meta *types.TokenCountMeta, info *relaycommon.RelayInfo) (int, error) {
if meta == nil {
return 0, errors.New("token count meta is nil")
}

if info.BillingMode == "" {
info.BillingMode = billing_setting.GetBillingMode(info.OriginModelName)
}
if info.BillingMode == billing_setting.BillingModeUTF8Bytes {
count := len(meta.CombineText)
common.SetContextKey(c, constant.ContextKeyPromptTokens, count)
return count, nil
}

// 是否统计token
if !constant.CountToken {
return 0, nil
}

if meta == nil {
return 0, errors.New("token count meta is nil")
}

if info.RelayFormat == types.RelayFormatOpenAIRealtime {
return 0, nil
}
Comment on lines +185 to 201

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win

🧩 Analysis chain

🏁 Script executed:

#!/bin/bash
# Description: Check how GetTokenCountMeta is implemented for realtime/openai-realtime request types.
rg -n -C5 'func.*GetTokenCountMeta' --type=go
rg -n -C5 'RelayFormatOpenAIRealtime' --type=go relay

Repository: QuantumNous/new-api

Length of output: 157


🏁 Script executed:

#!/bin/bash
set -euo pipefail

echo "== git status/stat =="
git status --short
git diff --stat || true

echo "== locate relevant files =="
fd -a 'token_counter.go|relay.go|types|openai.*realtime|billing_setting' . | sed 's#^\./##' | head -200

echo "== search key identifiers =="
rg -n -C4 'GetTokenCountMeta|RelayFormatOpenAIRealtime|BillingModeUTF8Bytes|CombineText|CountToken' . --glob '*.go' --glob '*.ts' --glob '*.tsx' | head -300

Repository: QuantumNous/new-api

Length of output: 20999


🏁 Script executed:

#!/bin/bash
set -euo pipefail

echo "== service/token_counter.go outline =="
ast-grep outline service/token_counter.go --view compact || true

echo "== service/token_counter.go relevant section =="
sed -n '160,230p' service/token_counter.go

echo "== billing_setting UTF8 bytes definitions/usages =="
rg -n -C4 'BillingModeUTF8Bytes|BilingMode|BillingMode' setting --glob '*.go' | head -200
fd -a '*.go' setting | xargs -r rg -n 'BillingModeUTF8Bytes|AddBillingMode|GetBillingMode|BillingMode' | head -200

echo "== relay valid request relevant section =="
sed -n '35,60p' relay/helper/valid_request.go

echo "== relay pricing relevant section =="
sed -n '80,115p' relay/helper/price.go

Repository: QuantumNous/new-api

Length of output: 7023


🏁 Script executed:

#!/bin/bash
set -euo pipefail

echo "== GetAndValidAudioRequest implementation =="
rg -n -C8 'GetAndValidAudioRequest|Build.*Audio' --glob '*.go' | head -220

echo "== realtime/audio token count metadata in OpenAI relay files =="
fd -a 'openai.*realtime|openai.*audio|audio' relaykit relay controller | sed 's#^\./##'
for f in $(fd -a 'openai.*realtime|openai.*audio|audio' relaykit relay controller); do
  echo "--- $f"
  rg -n -C6 'GetTokenCountMeta|Input|CombineText|openai_realtime|OpenAIRealtime|Audio' "$f" || true
done

echo "== pricing helper pre-consumption context =="
rg -n -C8 'SetEstimatePromptTokens|GetEstimatePromptTokens|PreConsumedTokens|promptTokens|BillingModeUTF8Bytes|RelayFormatOpenAIRealtime' relay --glob '*.go' | head -260

echo "== token_counter EstimateRequestToken continuation =="
sed -n '185,270p' service/token_counter.go

Repository: QuantumNous/new-api

Length of output: 202


🏁 Script executed:

#!/bin/bash
set -euo pipefail

echo "== files mentioning token count meta for realtime/audio =="
rg -n -C6 'RelayFormatOpenAIRealtime|RelayFormatOpenAIAudio|GetTokenCountMeta|CombineText|SetEstimatePromptTokens|GetEstimatePromptTokens' relay --glob '*.go' | head -400

echo "== list relevant relay dto types =="
fd -a 'openai.*realtime|realtime|openai.*audio|audio' relaykit relaykit/dto relay --print0 | xargs -0 -r -I{} sh -c 'echo "--- {}"; wc -l "{}" | sed "s#$# file:"; sed -n "1,220p" "{}"'

echo "== price go relevant =="
sed -n '85,120p' relay/helper/price.go

Repository: QuantumNous/new-api

Length of output: 50375


🏁 Script executed:

#!/bin/bash
set -euo pipefail

echo "== realtime DTO real file =="
sed -n '1,220p' relaykit/dto/realtime.go

echo "== realtime implementation relevant =="
sed -n '1,250p' relay/channel/openai/relay_realtime.go

echo "== pricing preConsume context =="
rg -n 'func preConsumeUsage|preConsumeUsage\(' relay --glob '*.go' -C8

echo "== realtime price/consume helpers =="
rg -n 'RealtimeUsage|audioToken|textToken|input_tokens|output_tokens|InputTokenDetails|Post.*Consume|quota' relay/helper relay channel/relay --glob '*.go' -C5 | head -400

echo "== relevant token counter section exact =="
sed -n '175,205p' service/token_counter.go

Repository: QuantumNous/new-api

Length of output: 41529


Preserve the realtime zero pre-consumer before UTF-8 byte billing.

RelayFormatOpenAIRealtime uses websocket usage accounting through preConsumeUsage; realtime request validation also creates a BaseRequest, so meta.CombineText is empty or not representative. Keep the info.RelayFormat == types.RelayFormatOpenAIRealtime zero-return before the BillingModeUTF8Bytes path, or guard it as a non-billable realtime case.

🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@service/token_counter.go` around lines 185 - 201, Move the
RelayFormatOpenAIRealtime zero-return check before the BillingModeUTF8Bytes
branch in the token-counting flow. Update the surrounding logic in the relevant
token-counting function so realtime requests return zero without billing based
on the empty or non-representative meta.CombineText, while preserving existing
behavior for other relay formats.

Expand Down
Loading