fix(usage): correct mapped model statistics

This commit is contained in:
visa2
2026-07-26 23:56:34 +08:00
parent 2730c1c43b
commit 1f45c99de7
7 changed files with 72 additions and 8 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,35 @@ 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_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: optionalNonEqualStringPtr(result.UpstreamModel, requestedModel),
ReasoningEffort: result.ReasoningEffort,
InboundEndpoint: optionalTrimmedStringPtr(input.InboundEndpoint),
UpstreamEndpoint: optionalTrimmedStringPtr(input.UpstreamEndpoint),
@@ -1335,6 +1335,40 @@ 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_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: optionalNonEqualStringPtr(result.UpstreamModel, requestedModel),
ServiceTier: result.ServiceTier,
ReasoningEffort: result.ReasoningEffort,
InboundEndpoint: optionalTrimmedStringPtr(input.InboundEndpoint),
@@ -11,8 +11,8 @@ func optionalTrimmedStringPtr(raw string) *string {
}
// 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.
// differs from compare; otherwise nil. Usage logging passes the requested
// model as compare so a channel mapping still records its effective upstream.
func optionalNonEqualStringPtr(value, compare string) *string {
if value == "" || value == compare {
return nil