Merge pull request #4933 from visa2/fix/usage-model-mapping-statistics

fix(usage): report channel-mapped requests under their real upstream model
This commit is contained in:
Wesley Liddick
2026-07-27 11:45:58 +08:00
committed by GitHub
7 changed files with 130 additions and 16 deletions
@@ -40,8 +40,8 @@ func TestResolveModelDimensionExpression(t *testing.T) {
want string
}{
{usagestats.ModelSourceRequested, "COALESCE(NULLIF(TRIM(requested_model), ''), model)"},
{usagestats.ModelSourceUpstream, "COALESCE(NULLIF(TRIM(upstream_model), ''), COALESCE(NULLIF(TRIM(requested_model), ''), model))"},
{usagestats.ModelSourceMapping, "(COALESCE(NULLIF(TRIM(requested_model), ''), model) || ' -> ' || COALESCE(NULLIF(TRIM(upstream_model), ''), COALESCE(NULLIF(TRIM(requested_model), ''), model)))"},
{usagestats.ModelSourceUpstream, "COALESCE(NULLIF(TRIM(upstream_model), ''), model)"},
{usagestats.ModelSourceMapping, "(COALESCE(NULLIF(TRIM(requested_model), ''), model) || ' -> ' || COALESCE(NULLIF(TRIM(upstream_model), ''), model))"},
{"", "COALESCE(NULLIF(TRIM(requested_model), ''), model)"},
{"invalid", "COALESCE(NULLIF(TRIM(requested_model), ''), model)"},
}
@@ -750,11 +750,12 @@ func resolveModelDimensionExpressionWithAlias(modelType, alias string) string {
return alias + "." + name
}
requestedExpr := fmt.Sprintf("COALESCE(NULLIF(TRIM(%s), ''), %s)", column("requested_model"), column("model"))
upstreamExpr := fmt.Sprintf("COALESCE(NULLIF(TRIM(%s), ''), %s)", column("upstream_model"), column("model"))
switch usagestats.NormalizeModelSource(modelType) {
case usagestats.ModelSourceUpstream:
return fmt.Sprintf("COALESCE(NULLIF(TRIM(%s), ''), %s)", column("upstream_model"), requestedExpr)
return upstreamExpr
case usagestats.ModelSourceMapping:
return fmt.Sprintf("(%s || ' -> ' || COALESCE(NULLIF(TRIM(%s), ''), %s))", requestedExpr, column("upstream_model"), requestedExpr)
return fmt.Sprintf("(%s || ' -> ' || %s)", requestedExpr, upstreamExpr)
default:
return requestedExpr
}
@@ -194,6 +194,64 @@ func TestGatewayServiceRecordUsage_PreservesRequestedAndUpstreamModels(t *testin
require.Equal(t, mappedModel, *usageRepo.lastLog.UpstreamModel)
}
func TestGatewayServiceRecordUsage_PreservesChannelMappedUpstreamModel(t *testing.T) {
usageRepo := &openAIRecordUsageLogRepoStub{inserted: true}
svc := newGatewayRecordUsageServiceForTest(usageRepo, &openAIRecordUsageUserRepoStub{}, &openAIRecordUsageSubRepoStub{})
err := svc.RecordUsage(context.Background(), &RecordUsageInput{
Result: &ForwardResult{
RequestID: "gateway_channel_mapping_models",
Usage: ClaudeUsage{InputTokens: 10, OutputTokens: 6},
Model: "gpt-5.6-terra",
UpstreamModel: "gpt-5.6-terra",
Duration: time.Second,
},
APIKey: &APIKey{ID: 501, Quota: 100},
User: &User{ID: 601},
Account: &Account{ID: 701},
ChannelUsageFields: ChannelUsageFields{
OriginalModel: "gpt-5.6-sol",
ChannelMappedModel: "gpt-5.6-terra",
},
})
require.NoError(t, err)
require.NotNil(t, usageRepo.lastLog)
require.Equal(t, "gpt-5.6-sol", usageRepo.lastLog.RequestedModel)
require.Equal(t, "gpt-5.6-terra", usageRepo.lastLog.Model)
require.NotNil(t, usageRepo.lastLog.UpstreamModel)
require.Equal(t, "gpt-5.6-terra", *usageRepo.lastLog.UpstreamModel)
}
func TestGatewayServiceRecordUsage_PreservesLoopedChannelAndAccountUpstreamModel(t *testing.T) {
usageRepo := &openAIRecordUsageLogRepoStub{inserted: true}
svc := newGatewayRecordUsageServiceForTest(usageRepo, &openAIRecordUsageUserRepoStub{}, &openAIRecordUsageSubRepoStub{})
err := svc.RecordUsage(context.Background(), &RecordUsageInput{
Result: &ForwardResult{
RequestID: "gateway_looped_mapping_models",
Usage: ClaudeUsage{InputTokens: 10, OutputTokens: 6},
Model: "gpt-5.6-terra",
UpstreamModel: "gpt-5.6-sol",
Duration: time.Second,
},
APIKey: &APIKey{ID: 501, Quota: 100},
User: &User{ID: 601},
Account: &Account{ID: 701},
ChannelUsageFields: ChannelUsageFields{
OriginalModel: "gpt-5.6-sol",
ChannelMappedModel: "gpt-5.6-terra",
},
})
require.NoError(t, err)
require.NotNil(t, usageRepo.lastLog)
require.Equal(t, "gpt-5.6-sol", usageRepo.lastLog.RequestedModel)
require.Equal(t, "gpt-5.6-terra", usageRepo.lastLog.Model)
require.NotNil(t, usageRepo.lastLog.UpstreamModel)
require.Equal(t, "gpt-5.6-sol", *usageRepo.lastLog.UpstreamModel)
}
func TestGatewayServiceRecordUsage_EmptyImageSizeDefaultsBeforeBillingAndPersistence(t *testing.T) {
imagePrice2K := 0.19
groupID := int64(901)
@@ -986,7 +986,7 @@ func (s *GatewayService) buildRecordUsageLog(
RequestID: requestID,
Model: result.Model,
RequestedModel: requestedModel,
UpstreamModel: optionalNonEqualStringPtr(result.UpstreamModel, result.Model),
UpstreamModel: optionalTrimmedStringPtr(result.UpstreamModel),
ReasoningEffort: result.ReasoningEffort,
InboundEndpoint: optionalTrimmedStringPtr(input.InboundEndpoint),
UpstreamEndpoint: optionalTrimmedStringPtr(input.UpstreamEndpoint),
@@ -1335,6 +1335,71 @@ func TestOpenAIGatewayServiceRecordUsage_UsesRequestedModelAndUpstreamModelMetad
require.Equal(t, 1, userRepo.deductCalls)
}
func TestOpenAIGatewayServiceRecordUsage_PreservesChannelMappedUpstreamModel(t *testing.T) {
usageRepo := &openAIRecordUsageLogRepoStub{inserted: true}
userRepo := &openAIRecordUsageUserRepoStub{}
subRepo := &openAIRecordUsageSubRepoStub{}
svc := newOpenAIRecordUsageServiceForTest(usageRepo, userRepo, subRepo, nil)
err := svc.RecordUsage(context.Background(), &OpenAIRecordUsageInput{
Result: &OpenAIForwardResult{
RequestID: "openai_channel_mapping_models",
Model: "gpt-5.6-terra",
UpstreamModel: "gpt-5.6-terra",
Usage: OpenAIUsage{
InputTokens: 20,
OutputTokens: 10,
},
Duration: time.Second,
},
APIKey: &APIKey{ID: 10},
User: &User{ID: 20},
Account: &Account{ID: 30},
ChannelUsageFields: ChannelUsageFields{
OriginalModel: "gpt-5.6-sol",
ChannelMappedModel: "gpt-5.6-terra",
},
})
require.NoError(t, err)
require.NotNil(t, usageRepo.lastLog)
require.Equal(t, "gpt-5.6-sol", usageRepo.lastLog.RequestedModel)
require.Equal(t, "gpt-5.6-terra", usageRepo.lastLog.Model)
require.NotNil(t, usageRepo.lastLog.UpstreamModel)
require.Equal(t, "gpt-5.6-terra", *usageRepo.lastLog.UpstreamModel)
}
func TestOpenAIGatewayServiceRecordUsage_PreservesLoopedChannelAndAccountUpstreamModel(t *testing.T) {
usageRepo := &openAIRecordUsageLogRepoStub{inserted: true}
userRepo := &openAIRecordUsageUserRepoStub{}
subRepo := &openAIRecordUsageSubRepoStub{}
svc := newOpenAIRecordUsageServiceForTest(usageRepo, userRepo, subRepo, nil)
err := svc.RecordUsage(context.Background(), &OpenAIRecordUsageInput{
Result: &OpenAIForwardResult{
RequestID: "openai_looped_mapping_models",
Model: "gpt-5.6-terra",
UpstreamModel: "gpt-5.6-sol",
Usage: OpenAIUsage{InputTokens: 20, OutputTokens: 10},
Duration: time.Second,
},
APIKey: &APIKey{ID: 10},
User: &User{ID: 20},
Account: &Account{ID: 30},
ChannelUsageFields: ChannelUsageFields{
OriginalModel: "gpt-5.6-sol",
ChannelMappedModel: "gpt-5.6-terra",
},
})
require.NoError(t, err)
require.NotNil(t, usageRepo.lastLog)
require.Equal(t, "gpt-5.6-sol", usageRepo.lastLog.RequestedModel)
require.Equal(t, "gpt-5.6-terra", usageRepo.lastLog.Model)
require.NotNil(t, usageRepo.lastLog.UpstreamModel)
require.Equal(t, "gpt-5.6-sol", *usageRepo.lastLog.UpstreamModel)
}
func TestOpenAIGatewayServiceRecordUsage_BillsMappedRequestsUsingRequestedModel(t *testing.T) {
usageRepo := &openAIRecordUsageLogRepoStub{inserted: true}
userRepo := &openAIRecordUsageUserRepoStub{}
@@ -256,7 +256,7 @@ func (s *OpenAIGatewayService) RecordUsage(ctx context.Context, input *OpenAIRec
RequestID: requestID,
Model: result.Model,
RequestedModel: requestedModel,
UpstreamModel: optionalNonEqualStringPtr(result.UpstreamModel, result.Model),
UpstreamModel: optionalTrimmedStringPtr(result.UpstreamModel),
ServiceTier: result.ServiceTier,
ReasoningEffort: result.ReasoningEffort,
InboundEndpoint: optionalTrimmedStringPtr(input.InboundEndpoint),
@@ -10,16 +10,6 @@ func optionalTrimmedStringPtr(raw string) *string {
return &trimmed
}
// optionalNonEqualStringPtr returns a pointer to value if it is non-empty and
// differs from compare; otherwise nil. Used to store upstream_model only when
// it differs from the requested model.
func optionalNonEqualStringPtr(value, compare string) *string {
if value == "" || value == compare {
return nil
}
return &value
}
func forwardResultBillingModel(requestedModel, upstreamModel string) string {
if trimmed := strings.TrimSpace(requestedModel); trimmed != "" {
return trimmed