mirror of
https://github.com/QuantumNous/new-api.git
synced 2026-09-11 14:41:21 +00:00
feat(relay): hosted-tool conversion fidelity, reasoning normalization, and billing usage integrity (#7137)
* feat(relaykit): preserve hosted tools across conversions - add protocol-neutral hosted-tool DTOs, conversion metadata, and loss policies - bridge citations, grounding metadata, and hosted-tool stream lifecycles - document the public conversion behavior and channel policy controls * refactor(relaykit): normalize reasoning and thinking intent - centralize provider-neutral reasoning intent, effort, and budget mappings - parse model suffixes at the host entry boundary while preserving provider-owned tails - keep adaptive Claude thinking and explicit zero-token compatibility consistent * fix(billing): preserve authoritative usage across relay hops - carry native BillingUsage sidecars through direct and streamed protocol bridges - merge partial and terminal usage monotonically with safe fallback settlement - retain cache metadata, penultimate usage, and per-call Gemini tool surcharges * feat(relay): bridge Responses with Claude and Gemini protocols - add direct request, response, and stream converters across supported relay formats - expose Claude count_tokens and Chat-to-Responses compatibility endpoints - carry conversion diagnostics through the host while retaining the curated public goldens * fix(relay): wire relaykit conversions into host channels - connect handlers, adaptors, and channel settings to the standalone conversion layer - keep model mapping, pricing identity, retries, and provider-specific suffix behavior aligned - ignore local audit artifacts and retain focused public regression coverage
This commit is contained in:
@@ -401,7 +401,7 @@ func NewBillingSession(c *gin.Context, relayInfo *relaycommon.RelayInfo, preCons
|
||||
funding: &SubscriptionFunding{
|
||||
requestId: relayInfo.RequestId,
|
||||
userId: relayInfo.UserId,
|
||||
modelName: relayInfo.OriginModelName,
|
||||
modelName: relayInfo.GetBillingModelName(),
|
||||
amount: subConsume,
|
||||
},
|
||||
}
|
||||
|
||||
+2
-156
@@ -1,10 +1,6 @@
|
||||
package service
|
||||
|
||||
import (
|
||||
"strings"
|
||||
|
||||
"github.com/QuantumNous/new-api/relaykit/dto"
|
||||
)
|
||||
import "github.com/QuantumNous/new-api/relaykit/dto"
|
||||
|
||||
const (
|
||||
usageBillingPathLocal = "local"
|
||||
@@ -70,155 +66,5 @@ func usageFromBillingUsage(usage *dto.Usage) (*dto.Usage, bool) {
|
||||
if usage == nil || usage.BillingUsage == nil {
|
||||
return nil, false
|
||||
}
|
||||
billingUsage := usage.BillingUsage
|
||||
source := strings.TrimSpace(billingUsage.Source)
|
||||
semantic := strings.TrimSpace(billingUsage.Semantic)
|
||||
|
||||
if billingUsage.OpenAIUsage != nil &&
|
||||
(strings.EqualFold(source, dto.BillingUsageSourceOAIChat) ||
|
||||
strings.EqualFold(source, dto.BillingUsageSourceOAIResponses) ||
|
||||
strings.EqualFold(semantic, dto.BillingUsageSemanticOpenAI)) {
|
||||
return usageFromOpenAIBillingUsage(billingUsage), true
|
||||
}
|
||||
|
||||
if billingUsage.ClaudeUsage != nil &&
|
||||
(strings.EqualFold(source, dto.BillingUsageSourceClaudeMessages) ||
|
||||
strings.EqualFold(semantic, dto.BillingUsageSemanticAnthropic)) {
|
||||
return usageFromClaudeBillingUsage(billingUsage), true
|
||||
}
|
||||
|
||||
if billingUsage.GeminiUsageMetadata != nil &&
|
||||
(strings.EqualFold(source, dto.BillingUsageSourceGeminiChat) ||
|
||||
strings.EqualFold(semantic, dto.BillingUsageSemanticGemini)) {
|
||||
return usageFromGeminiBillingUsage(billingUsage), true
|
||||
}
|
||||
|
||||
return nil, false
|
||||
}
|
||||
|
||||
func usageFromOpenAIBillingUsage(billingUsage *dto.BillingUsage) *dto.Usage {
|
||||
usage := *billingUsage.OpenAIUsage
|
||||
if usage.PromptTokens == 0 && usage.InputTokens > 0 {
|
||||
usage.PromptTokens = usage.InputTokens
|
||||
}
|
||||
if usage.CompletionTokens == 0 && usage.OutputTokens > 0 {
|
||||
usage.CompletionTokens = usage.OutputTokens
|
||||
}
|
||||
if usage.InputTokens == 0 && usage.PromptTokens > 0 {
|
||||
usage.InputTokens = usage.PromptTokens
|
||||
}
|
||||
if usage.OutputTokens == 0 && usage.CompletionTokens > 0 {
|
||||
usage.OutputTokens = usage.CompletionTokens
|
||||
}
|
||||
if usage.TotalTokens == 0 {
|
||||
usage.TotalTokens = usage.PromptTokens + usage.CompletionTokens
|
||||
}
|
||||
if inputDetails := usage.InputTokensDetails; inputDetails != nil {
|
||||
if usage.PromptTokensDetails.CachedTokens == 0 && inputDetails.CachedTokens > 0 {
|
||||
usage.PromptTokensDetails.CachedTokens = inputDetails.CachedTokens
|
||||
}
|
||||
if usage.PromptTokensDetails.CachedCreationTokens == 0 && inputDetails.CachedCreationTokens > 0 {
|
||||
usage.PromptTokensDetails.CachedCreationTokens = inputDetails.CachedCreationTokens
|
||||
}
|
||||
if usage.PromptTokensDetails.CacheWriteTokens == 0 && inputDetails.CacheWriteTokens > 0 {
|
||||
usage.PromptTokensDetails.CacheWriteTokens = inputDetails.CacheWriteTokens
|
||||
}
|
||||
if usage.PromptTokensDetails.TextTokens == 0 && inputDetails.TextTokens > 0 {
|
||||
usage.PromptTokensDetails.TextTokens = inputDetails.TextTokens
|
||||
}
|
||||
if usage.PromptTokensDetails.ImageTokens == 0 && inputDetails.ImageTokens > 0 {
|
||||
usage.PromptTokensDetails.ImageTokens = inputDetails.ImageTokens
|
||||
}
|
||||
if usage.PromptTokensDetails.AudioTokens == 0 && inputDetails.AudioTokens > 0 {
|
||||
usage.PromptTokensDetails.AudioTokens = inputDetails.AudioTokens
|
||||
}
|
||||
}
|
||||
if usage.PromptTokensDetails.CachedTokens == 0 && usage.PromptCacheHitTokens > 0 {
|
||||
usage.PromptTokensDetails.CachedTokens = usage.PromptCacheHitTokens
|
||||
}
|
||||
usage.UsageSemantic = dto.BillingUsageSemanticOpenAI
|
||||
usage.UsageSource = billingUsage.Source
|
||||
usage.BillingUsage = dto.CloneBillingUsage(billingUsage)
|
||||
return &usage
|
||||
}
|
||||
|
||||
func usageFromClaudeBillingUsage(billingUsage *dto.BillingUsage) *dto.Usage {
|
||||
claudeUsage := billingUsage.ClaudeUsage
|
||||
cacheCreation5m := claudeUsage.GetCacheCreation5mTokens()
|
||||
if cacheCreation5m == 0 {
|
||||
cacheCreation5m = claudeUsage.ClaudeCacheCreation5mTokens
|
||||
}
|
||||
cacheCreation1h := claudeUsage.GetCacheCreation1hTokens()
|
||||
if cacheCreation1h == 0 {
|
||||
cacheCreation1h = claudeUsage.ClaudeCacheCreation1hTokens
|
||||
}
|
||||
|
||||
usage := &dto.Usage{
|
||||
PromptTokens: claudeUsage.InputTokens,
|
||||
CompletionTokens: claudeUsage.OutputTokens,
|
||||
TotalTokens: claudeUsage.InputTokens + claudeUsage.OutputTokens,
|
||||
InputTokens: claudeUsage.InputTokens + claudeUsage.CacheReadInputTokens + claudeUsage.CacheCreationInputTokens,
|
||||
OutputTokens: claudeUsage.OutputTokens,
|
||||
UsageSemantic: dto.BillingUsageSemanticAnthropic,
|
||||
UsageSource: dto.BillingUsageSourceClaudeMessages,
|
||||
BillingUsage: dto.CloneBillingUsage(billingUsage),
|
||||
ClaudeCacheCreation5mTokens: cacheCreation5m,
|
||||
ClaudeCacheCreation1hTokens: cacheCreation1h,
|
||||
}
|
||||
usage.PromptTokensDetails.CachedTokens = claudeUsage.CacheReadInputTokens
|
||||
usage.PromptTokensDetails.CachedCreationTokens = claudeUsage.CacheCreationInputTokens
|
||||
return usage
|
||||
}
|
||||
|
||||
func usageFromGeminiBillingUsage(billingUsage *dto.BillingUsage) *dto.Usage {
|
||||
metadata := *billingUsage.GeminiUsageMetadata
|
||||
promptTokens := metadata.PromptTokenCount + metadata.ToolUsePromptTokenCount
|
||||
usage := &dto.Usage{
|
||||
PromptTokens: promptTokens,
|
||||
CompletionTokens: metadata.CandidatesTokenCount + metadata.ThoughtsTokenCount,
|
||||
TotalTokens: metadata.TotalTokenCount,
|
||||
UsageSemantic: dto.BillingUsageSemanticGemini,
|
||||
UsageSource: dto.BillingUsageSourceGeminiChat,
|
||||
BillingUsage: dto.CloneBillingUsage(billingUsage),
|
||||
}
|
||||
usage.CompletionTokenDetails.ReasoningTokens = metadata.ThoughtsTokenCount
|
||||
usage.PromptTokensDetails.CachedTokens = metadata.CachedContentTokenCount
|
||||
|
||||
for _, detail := range metadata.PromptTokensDetails {
|
||||
addGeminiInputTokenDetail(&usage.PromptTokensDetails, detail)
|
||||
}
|
||||
for _, detail := range metadata.ToolUsePromptTokensDetails {
|
||||
addGeminiInputTokenDetail(&usage.PromptTokensDetails, detail)
|
||||
}
|
||||
for _, detail := range metadata.CandidatesTokensDetails {
|
||||
switch detail.Modality {
|
||||
case "IMAGE":
|
||||
usage.CompletionTokenDetails.ImageTokens += detail.TokenCount
|
||||
case "AUDIO":
|
||||
usage.CompletionTokenDetails.AudioTokens += detail.TokenCount
|
||||
case "TEXT":
|
||||
usage.CompletionTokenDetails.TextTokens += detail.TokenCount
|
||||
}
|
||||
}
|
||||
|
||||
if usage.TotalTokens == 0 {
|
||||
usage.TotalTokens = usage.PromptTokens + usage.CompletionTokens
|
||||
} else if usage.CompletionTokens <= 0 {
|
||||
usage.CompletionTokens = usage.TotalTokens - usage.PromptTokens
|
||||
}
|
||||
if usage.PromptTokens > 0 && usage.PromptTokensDetails.TextTokens == 0 && usage.PromptTokensDetails.AudioTokens == 0 {
|
||||
usage.PromptTokensDetails.TextTokens = usage.PromptTokens
|
||||
}
|
||||
return usage
|
||||
}
|
||||
|
||||
func addGeminiInputTokenDetail(details *dto.InputTokenDetails, detail dto.GeminiPromptTokensDetails) {
|
||||
switch detail.Modality {
|
||||
case "AUDIO":
|
||||
details.AudioTokens += detail.TokenCount
|
||||
case "IMAGE":
|
||||
details.ImageTokens += detail.TokenCount
|
||||
case "TEXT":
|
||||
details.TextTokens += detail.TokenCount
|
||||
}
|
||||
return usage.BillingUsage.CanonicalUsage()
|
||||
}
|
||||
|
||||
@@ -47,7 +47,7 @@ func attachQuotaSaturation(ctx *gin.Context, relayInfo *relaycommon.RelayInfo, o
|
||||
}
|
||||
attachQuotaSaturationToOther(other, clamp)
|
||||
logger.LogWarn(ctx, fmt.Sprintf("quota saturation on consume log: op=%s kind=%s original=%g clamped=%d user=%d model=%s",
|
||||
clamp.Op, clamp.Kind, clamp.Original, clamp.Clamped, relayInfo.UserId, relayInfo.OriginModelName))
|
||||
clamp.Op, clamp.Kind, clamp.Original, clamp.Clamped, relayInfo.UserId, relayInfo.GetBillingModelName()))
|
||||
}
|
||||
|
||||
func appendRequestPath(ctx *gin.Context, relayInfo *relaycommon.RelayInfo, other map[string]interface{}) {
|
||||
@@ -95,6 +95,15 @@ func GenerateTextOtherInfo(ctx *gin.Context, relayInfo *relaycommon.RelayInfo, m
|
||||
|
||||
adminInfo := make(map[string]interface{})
|
||||
adminInfo["use_channel"] = ctx.GetStringSlice("use_channel")
|
||||
if billingModel := relayInfo.GetBillingModelName(); billingModel != "" && billingModel != relayInfo.OriginModelName {
|
||||
adminInfo["billing_model"] = billingModel
|
||||
}
|
||||
if diagnostics := relayInfo.ConversionDiagnostics(); len(diagnostics) > 0 {
|
||||
adminInfo["conversion_diagnostics"] = diagnostics
|
||||
}
|
||||
if relayInfo.ConversionDiagnosticsTruncated() {
|
||||
adminInfo["conversion_diagnostics_truncated"] = true
|
||||
}
|
||||
isMultiKey := common.GetContextKeyBool(ctx, constant.ContextKeyChannelIsMultiKey)
|
||||
if isMultiKey {
|
||||
adminInfo["is_multi_key"] = true
|
||||
|
||||
+7
-6
@@ -303,9 +303,10 @@ func PostAudioConsumeQuota(ctx *gin.Context, relayInfo *relaycommon.RelayInfo, u
|
||||
audioOutTokens := usage.CompletionTokenDetails.AudioTokens
|
||||
|
||||
tokenName := ctx.GetString("token_name")
|
||||
completionRatio := decimal.NewFromFloat(ratio_setting.GetCompletionRatio(relayInfo.OriginModelName))
|
||||
audioRatio := decimal.NewFromFloat(ratio_setting.GetAudioRatio(relayInfo.OriginModelName))
|
||||
audioCompletionRatio := decimal.NewFromFloat(ratio_setting.GetAudioCompletionRatio(relayInfo.OriginModelName))
|
||||
billingModelName := relayInfo.GetBillingModelName()
|
||||
completionRatio := decimal.NewFromFloat(ratio_setting.GetCompletionRatio(billingModelName))
|
||||
audioRatio := decimal.NewFromFloat(ratio_setting.GetAudioRatio(billingModelName))
|
||||
audioCompletionRatio := decimal.NewFromFloat(ratio_setting.GetAudioCompletionRatio(billingModelName))
|
||||
|
||||
modelRatio := relayInfo.PriceData.ModelRatio
|
||||
groupRatio := relayInfo.PriceData.GroupRatioInfo.GroupRatio
|
||||
@@ -321,7 +322,7 @@ func PostAudioConsumeQuota(ctx *gin.Context, relayInfo *relaycommon.RelayInfo, u
|
||||
TextTokens: textOutTokens,
|
||||
AudioTokens: audioOutTokens,
|
||||
},
|
||||
ModelName: relayInfo.OriginModelName,
|
||||
ModelName: billingModelName,
|
||||
UsePrice: usePrice,
|
||||
ModelRatio: modelRatio,
|
||||
GroupRatio: groupRatio,
|
||||
@@ -349,7 +350,7 @@ func PostAudioConsumeQuota(ctx *gin.Context, relayInfo *relaycommon.RelayInfo, u
|
||||
quota = 0
|
||||
logContent += "(可能是上游超时)"
|
||||
logger.LogError(ctx, fmt.Sprintf("total tokens is 0, cannot consume quota, userId %d, channelId %d, "+
|
||||
"tokenId %d, model %s, pre-consumed quota %d", relayInfo.UserId, relayInfo.ChannelId, relayInfo.TokenId, relayInfo.OriginModelName, relayInfo.FinalPreConsumedQuota))
|
||||
"tokenId %d, model %s, pre-consumed quota %d", relayInfo.UserId, relayInfo.ChannelId, relayInfo.TokenId, billingModelName, relayInfo.FinalPreConsumedQuota))
|
||||
} else {
|
||||
model.UpdateUserUsedQuotaAndRequestCount(relayInfo.UserId, quota)
|
||||
model.UpdateChannelUsedQuota(relayInfo.ChannelId, quota)
|
||||
@@ -359,7 +360,7 @@ func PostAudioConsumeQuota(ctx *gin.Context, relayInfo *relaycommon.RelayInfo, u
|
||||
logger.LogError(ctx, "error settling billing: "+err.Error())
|
||||
}
|
||||
|
||||
logModel := relayInfo.OriginModelName
|
||||
logModel := billingModelName
|
||||
if extraContent != "" {
|
||||
logContent += ", " + extraContent
|
||||
}
|
||||
|
||||
@@ -24,15 +24,27 @@ func init() {
|
||||
}
|
||||
|
||||
func ConvertRequest(c *gin.Context, info *relaycommon.RelayInfo, target types.RelayFormat, request any) (*relayconvert.RequestResult, error) {
|
||||
return relayconvert.ConvertRequest(c, info, target, request)
|
||||
result, err := relayconvert.ConvertRequest(c, info, target, request)
|
||||
if result != nil {
|
||||
info.RecordConversionDiagnostics(c, result.Diagnostics)
|
||||
}
|
||||
return result, err
|
||||
}
|
||||
|
||||
func ConvertRequestByID(c *gin.Context, info *relaycommon.RelayInfo, converter string, request any) (*relayconvert.RequestResult, error) {
|
||||
return relayconvert.ConvertRequestByID(c, info, converter, request)
|
||||
result, err := relayconvert.ConvertRequestByID(c, info, converter, request)
|
||||
if result != nil {
|
||||
info.RecordConversionDiagnostics(c, result.Diagnostics)
|
||||
}
|
||||
return result, err
|
||||
}
|
||||
|
||||
func ConvertRequestVia(c *gin.Context, info *relaycommon.RelayInfo, request any, path ...types.RelayFormat) (*relayconvert.RequestResult, error) {
|
||||
return relayconvert.ConvertRequestVia(c, info, request, path...)
|
||||
result, err := relayconvert.ConvertRequestVia(c, info, request, path...)
|
||||
if result != nil {
|
||||
info.RecordConversionDiagnostics(c, result.Diagnostics)
|
||||
}
|
||||
return result, err
|
||||
}
|
||||
|
||||
func ClaudeToOpenAIRequest(claudeRequest dto.ClaudeRequest, info *relaycommon.RelayInfo) (*dto.GeneralOpenAIRequest, error) {
|
||||
|
||||
@@ -0,0 +1,40 @@
|
||||
package service
|
||||
|
||||
import (
|
||||
relaycommon "github.com/QuantumNous/new-api/relay/common"
|
||||
"github.com/QuantumNous/new-api/relaykit/relayconvert"
|
||||
"github.com/QuantumNous/new-api/relaykit/types"
|
||||
"github.com/gin-gonic/gin"
|
||||
)
|
||||
|
||||
func ConvertResponse(c *gin.Context, info *relaycommon.RelayInfo, target types.RelayFormat, response any) (*relayconvert.ResponseResult, error) {
|
||||
result, err := relayconvert.ConvertResponse(c, info, target, response)
|
||||
if result != nil {
|
||||
info.RecordConversionDiagnostics(c, result.Diagnostics)
|
||||
}
|
||||
return result, err
|
||||
}
|
||||
|
||||
func ConvertStreamResponse(c *gin.Context, info *relaycommon.RelayInfo, target types.RelayFormat, response any) (*relayconvert.ResponseResult, error) {
|
||||
result, err := relayconvert.ConvertStreamResponse(c, info, target, response)
|
||||
if result != nil {
|
||||
info.RecordConversionDiagnostics(c, result.Diagnostics)
|
||||
}
|
||||
return result, err
|
||||
}
|
||||
|
||||
func ConvertStreamResponseChunk(c *gin.Context, info *relaycommon.RelayInfo, state *relayconvert.ResponseStreamState, response any) ([]relayconvert.ResponseResult, error) {
|
||||
results, err := relayconvert.ConvertStreamResponseChunk(c, info, state, response)
|
||||
if state != nil {
|
||||
info.RecordConversionDiagnostics(c, state.Diagnostics())
|
||||
}
|
||||
return results, err
|
||||
}
|
||||
|
||||
func FinalizeStreamResponse(c *gin.Context, info *relaycommon.RelayInfo, state *relayconvert.ResponseStreamState) ([]relayconvert.ResponseResult, error) {
|
||||
results, err := relayconvert.FinalizeStreamResponse(c, info, state)
|
||||
if state != nil {
|
||||
info.RecordConversionDiagnostics(c, state.Diagnostics())
|
||||
}
|
||||
return results, err
|
||||
}
|
||||
@@ -230,7 +230,7 @@ func composeTieredTextQuota(relayInfo *relaycommon.RelayInfo, summary textQuotaS
|
||||
// the result with tiered billing, affinity observation and logging.
|
||||
func calculateTextQuotaSummary(ctx *gin.Context, relayInfo *relaycommon.RelayInfo, usage *dto.Usage) textQuotaSummary {
|
||||
summary := textQuotaSummary{
|
||||
ModelName: relayInfo.OriginModelName,
|
||||
ModelName: relayInfo.GetBillingModelName(),
|
||||
TokenName: ctx.GetString("token_name"),
|
||||
UseTimeSeconds: time.Now().Unix() - relayInfo.StartTime.Unix(),
|
||||
CompletionRatio: relayInfo.PriceData.CompletionRatio,
|
||||
|
||||
@@ -369,6 +369,53 @@ func TestUsageFromOpenAIBillingUsageFallsBackToPromptCacheHitTokens(t *testing.T
|
||||
require.Equal(t, 35, usage.PromptTokensDetails.CachedTokens)
|
||||
}
|
||||
|
||||
func TestCalculateTextQuotaSummaryNormalizesOpenAIResponsesBillingUsageDetails(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
w := httptest.NewRecorder()
|
||||
ctx, _ := gin.CreateTestContext(w)
|
||||
|
||||
relayInfo := &relaycommon.RelayInfo{
|
||||
RelayFormat: types.RelayFormatClaude,
|
||||
OriginModelName: "gpt-5.6-sol",
|
||||
PriceData: hosttypes.PriceData{
|
||||
ModelRatio: 1,
|
||||
CompletionRatio: 2,
|
||||
CacheRatio: 0.5,
|
||||
CacheCreationRatio: 2,
|
||||
GroupRatioInfo: hosttypes.GroupRatioInfo{GroupRatio: 1},
|
||||
},
|
||||
StartTime: time.Now(),
|
||||
}
|
||||
|
||||
responsesDetails := dto.InputTokenDetails{
|
||||
CachedTokens: 80,
|
||||
CacheWriteTokens: 10,
|
||||
TextTokens: 100,
|
||||
}
|
||||
usage := &dto.Usage{
|
||||
PromptTokens: 999,
|
||||
CompletionTokens: 999,
|
||||
BillingUsage: dto.NewOpenAIResponsesBillingUsage(&dto.Usage{
|
||||
InputTokens: 100,
|
||||
OutputTokens: 10,
|
||||
TotalTokens: 110,
|
||||
InputTokensDetails: &responsesDetails,
|
||||
}),
|
||||
}
|
||||
|
||||
effectiveUsage := effectiveBillingUsage(usage)
|
||||
summary := calculateTextQuotaSummary(ctx, relayInfo, effectiveUsage)
|
||||
|
||||
require.Equal(t, dto.BillingUsageSourceOAIResponses, effectiveUsage.UsageSource)
|
||||
require.Equal(t, responsesDetails, effectiveUsage.PromptTokensDetails)
|
||||
require.Equal(t, 100, summary.PromptTokens)
|
||||
require.Equal(t, 10, summary.CompletionTokens)
|
||||
require.Equal(t, 80, summary.CacheTokens)
|
||||
require.Equal(t, 10, summary.CacheCreationTokens)
|
||||
// (100-80-10) + 80*0.5 + 10*2 + 10*2 = 90
|
||||
require.Equal(t, 90, summary.Quota)
|
||||
}
|
||||
|
||||
func TestUsageBillingPathForLog(t *testing.T) {
|
||||
require.Equal(t, usageBillingPathAnthropic, usageBillingPathForLog(true, &dto.Usage{
|
||||
BillingUsage: dto.NewClaudeMessagesBillingUsage(&dto.ClaudeUsage{InputTokens: 1}),
|
||||
@@ -1044,6 +1091,38 @@ func TestCalculateTextToolCallSurchargeGeminiGoogleSearch(t *testing.T) {
|
||||
assert.Equal(t, 14.0, summary.ToolSurchargeItems[0].Price)
|
||||
}
|
||||
|
||||
func TestCalculateTextToolCallSurchargeGeminiFunctionCall(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
ctx, _ := gin.CreateTestContext(httptest.NewRecorder())
|
||||
|
||||
operation_setting.SetToolPriceForTest("gemini_surcharge_fn", 5.0)
|
||||
t.Cleanup(func() {
|
||||
operation_setting.DeleteToolPriceForTest("gemini_surcharge_fn")
|
||||
})
|
||||
|
||||
relayInfo := &relaycommon.RelayInfo{
|
||||
OriginModelName: "gemini-2.5-flash",
|
||||
ResponsesUsageInfo: &relaycommon.ResponsesUsageInfo{
|
||||
BuiltInTools: map[string]*relaycommon.BuildInToolInfo{
|
||||
"gemini_surcharge_fn": {CallCount: 2},
|
||||
},
|
||||
},
|
||||
}
|
||||
summary := &textQuotaSummary{ModelName: "gemini-2.5-flash", GroupRatio: 1}
|
||||
|
||||
surcharge := calculateTextToolCallSurcharge(ctx, relayInfo, summary)
|
||||
expected := decimal.NewFromFloat(5.0 * 2 / 1000).Mul(decimal.NewFromFloat(common.QuotaPerUnit))
|
||||
assert.True(t, expected.Equal(surcharge), "got %s want %s", surcharge, expected)
|
||||
require.Len(t, summary.ToolSurchargeItems, 1)
|
||||
assert.Equal(t, "gemini_surcharge_fn", summary.ToolSurchargeItems[0].Name)
|
||||
assert.Equal(t, 2, summary.ToolSurchargeItems[0].Count)
|
||||
assert.Equal(t, 5.0, summary.ToolSurchargeItems[0].Price)
|
||||
|
||||
other := map[string]interface{}{}
|
||||
appendToolSurchargeLogInfo(other, summary.ToolSurchargeItems)
|
||||
assert.Equal(t, summary.ToolSurchargeItems, other["tool_surcharges"])
|
||||
}
|
||||
|
||||
func TestCalculateTextToolCallSurchargeImageGenerationDefaultPrice(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
ctx, _ := gin.CreateTestContext(httptest.NewRecorder())
|
||||
|
||||
@@ -181,7 +181,13 @@ func EstimateRequestToken(c *gin.Context, meta *types.TokenCountMeta, info *rela
|
||||
if !constant.CountToken {
|
||||
return 0, nil
|
||||
}
|
||||
return CountRequestToken(c, meta, info)
|
||||
}
|
||||
|
||||
// CountRequestToken counts request tokens regardless of the billing estimation
|
||||
// switch. Utility endpoints such as Claude's messages/count_tokens must remain
|
||||
// available even when operators disable request-token estimation for relays.
|
||||
func CountRequestToken(c *gin.Context, meta *types.TokenCountMeta, info *relaycommon.RelayInfo) (int, error) {
|
||||
if meta == nil {
|
||||
return 0, errors.New("token count meta is nil")
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user