From dc1d814505a51f64988de968ba09677801580900 Mon Sep 17 00:00:00 2001 From: feitianbubu Date: Thu, 9 Jul 2026 13:44:30 +0800 Subject: [PATCH 1/2] fix: bill realtime transcription usage returned on completed events --- dto/realtime.go | 3 +++ relay/channel/openai/relay_realtime.go | 23 +++++++++++++++++++++++ 2 files changed, 26 insertions(+) diff --git a/dto/realtime.go b/dto/realtime.go index 0fbfb86f0d7c..1fba6967f8de 100644 --- a/dto/realtime.go +++ b/dto/realtime.go @@ -19,6 +19,7 @@ const ( RealtimeEventResponseFunctionCallArgumentsDelta = "response.function_call_arguments.delta" RealtimeEventResponseFunctionCallArgumentsDone = "response.function_call_arguments.done" RealtimeEventConversationItemCreated = "conversation.item.created" + RealtimeEventInputAudioTranscriptionCompleted = "conversation.item.input_audio_transcription.completed" ) type RealtimeEvent struct { @@ -31,6 +32,8 @@ type RealtimeEvent struct { Response *RealtimeResponse `json:"response,omitempty"` Delta string `json:"delta,omitempty"` Audio string `json:"audio,omitempty"` + // GA 转写结果事件(input_audio_transcription.completed)自带转写模型的 usage + Usage *RealtimeUsage `json:"usage,omitempty"` } type RealtimeResponse struct { diff --git a/relay/channel/openai/relay_realtime.go b/relay/channel/openai/relay_realtime.go index bb5c3587f886..add665d3e6a3 100644 --- a/relay/channel/openai/relay_realtime.go +++ b/relay/channel/openai/relay_realtime.go @@ -167,6 +167,29 @@ func OpenaiRealtimeHandler(c *gin.Context, info *relaycommon.RelayInfo) (*types. logger.LogInfo(c, fmt.Sprintf("realtime streaming localUsage: %v", localUsage)) logger.LogInfo(c, fmt.Sprintf("realtime streaming localUsage: %v", localUsage)) + } else if realtimeEvent.Type == dto.RealtimeEventInputAudioTranscriptionCompleted && realtimeEvent.Usage != nil && realtimeEvent.Usage.TotalTokens > 0 { + // GA 转写会话没有 response.done,转写模型的官方 usage 附在 completed 事件上,与 response.done 同等计费; + // whisper 系按时长返回(token 全 0)时不走此分支,仍用本地估算兜底 + usage.TotalTokens += realtimeEvent.Usage.TotalTokens + usage.InputTokens += realtimeEvent.Usage.InputTokens + usage.OutputTokens += realtimeEvent.Usage.OutputTokens + usage.InputTokenDetails.AudioTokens += realtimeEvent.Usage.InputTokenDetails.AudioTokens + usage.InputTokenDetails.CachedTokens += realtimeEvent.Usage.InputTokenDetails.CachedTokens + usage.InputTokenDetails.TextTokens += realtimeEvent.Usage.InputTokenDetails.TextTokens + usage.OutputTokenDetails.AudioTokens += realtimeEvent.Usage.OutputTokenDetails.AudioTokens + // GA 转写 usage 不带 output_token_details,而计费公式只认明细字段,明细缺失时输出全按文本补记 + outTextTokens := realtimeEvent.Usage.OutputTokenDetails.TextTokens + if outTextTokens == 0 && realtimeEvent.Usage.OutputTokenDetails.AudioTokens == 0 { + outTextTokens = realtimeEvent.Usage.OutputTokens + } + usage.OutputTokenDetails.TextTokens += outTextTokens + if err := preConsumeUsage(c, info, usage, sumUsage); err != nil { + errChan <- fmt.Errorf("error consume usage: %v", err) + return + } + // 官方 usage 已入账,清掉本句攒下的本地估算,避免重复 + usage = &dto.RealtimeUsage{} + localUsage = &dto.RealtimeUsage{} } else if realtimeEvent.Type == dto.RealtimeEventTypeSessionUpdated || realtimeEvent.Type == dto.RealtimeEventTypeSessionCreated { realtimeSession := realtimeEvent.Session if realtimeSession != nil { From 9b835ac03aa2c97367f3cc63ef45a8230b189d1e Mon Sep 17 00:00:00 2001 From: CaIon Date: Mon, 20 Jul 2026 21:00:01 +0800 Subject: [PATCH 2/2] fix: bill realtime transcription with ASR model --- relay/channel/openai/relay_realtime.go | 114 +++++--- relay/channel/openai/relay_realtime_test.go | 297 ++++++++++++++++++++ relay/common/realtime_transcription.go | 67 +++++ relay/common/realtime_transcription_test.go | 40 +++ relay/common/relay_info.go | 4 + service/quota.go | 155 +++++++--- service/quota_realtime_test.go | 63 +++++ 7 files changed, 664 insertions(+), 76 deletions(-) create mode 100644 relay/channel/openai/relay_realtime_test.go create mode 100644 relay/common/realtime_transcription.go create mode 100644 relay/common/realtime_transcription_test.go create mode 100644 service/quota_realtime_test.go diff --git a/relay/channel/openai/relay_realtime.go b/relay/channel/openai/relay_realtime.go index add665d3e6a3..a352756a047d 100644 --- a/relay/channel/openai/relay_realtime.go +++ b/relay/channel/openai/relay_realtime.go @@ -22,6 +22,7 @@ func OpenaiRealtimeHandler(c *gin.Context, info *relaycommon.RelayInfo) (*types. } info.IsStream = true + info.InitRealtimeTranscriptionState() clientConn := info.ClientWs targetConn := info.TargetWs @@ -34,6 +35,7 @@ func OpenaiRealtimeHandler(c *gin.Context, info *relaycommon.RelayInfo) (*types. usage := &dto.RealtimeUsage{} localUsage := &dto.RealtimeUsage{} sumUsage := &dto.RealtimeUsage{} + transcriptionSumUsage := &dto.RealtimeUsage{} gopool.Go(func() { defer func() { @@ -67,6 +69,7 @@ func OpenaiRealtimeHandler(c *gin.Context, info *relaycommon.RelayInfo) (*types. if realtimeEvent.Session.Tools != nil { info.RealtimeTools = realtimeEvent.Session.Tools } + captureRealtimeTranscriptionModel(info, realtimeEvent.Session) } } @@ -125,15 +128,8 @@ func OpenaiRealtimeHandler(c *gin.Context, info *relaycommon.RelayInfo) (*types. if realtimeEvent.Type == dto.RealtimeEventTypeResponseDone { realtimeUsage := realtimeEvent.Response.Usage if realtimeUsage != nil { - usage.TotalTokens += realtimeUsage.TotalTokens - usage.InputTokens += realtimeUsage.InputTokens - usage.OutputTokens += realtimeUsage.OutputTokens - usage.InputTokenDetails.AudioTokens += realtimeUsage.InputTokenDetails.AudioTokens - usage.InputTokenDetails.CachedTokens += realtimeUsage.InputTokenDetails.CachedTokens - usage.InputTokenDetails.TextTokens += realtimeUsage.InputTokenDetails.TextTokens - usage.OutputTokenDetails.AudioTokens += realtimeUsage.OutputTokenDetails.AudioTokens - usage.OutputTokenDetails.TextTokens += realtimeUsage.OutputTokenDetails.TextTokens - err := preConsumeUsage(c, info, usage, sumUsage) + *usage = addRealtimeUsage(*usage, *realtimeUsage) + _, err := preConsumeUsage(c, info, info.OriginModelName, usage, sumUsage) if err != nil { errChan <- fmt.Errorf("error consume usage: %v", err) return @@ -154,7 +150,7 @@ func OpenaiRealtimeHandler(c *gin.Context, info *relaycommon.RelayInfo) (*types. localUsage.InputTokens += textToken + audioToken localUsage.InputTokenDetails.TextTokens += textToken localUsage.InputTokenDetails.AudioTokens += audioToken - err = preConsumeUsage(c, info, localUsage, sumUsage) + _, err = preConsumeUsage(c, info, info.OriginModelName, localUsage, sumUsage) if err != nil { errChan <- fmt.Errorf("error consume usage: %v", err) return @@ -167,35 +163,25 @@ func OpenaiRealtimeHandler(c *gin.Context, info *relaycommon.RelayInfo) (*types. logger.LogInfo(c, fmt.Sprintf("realtime streaming localUsage: %v", localUsage)) logger.LogInfo(c, fmt.Sprintf("realtime streaming localUsage: %v", localUsage)) - } else if realtimeEvent.Type == dto.RealtimeEventInputAudioTranscriptionCompleted && realtimeEvent.Usage != nil && realtimeEvent.Usage.TotalTokens > 0 { + } else if transcriptionBilling, ok := realtimeTranscriptionBilling(info.OriginModelName, info.GetRealtimeTranscriptionModel(), realtimeEvent); ok { // GA 转写会话没有 response.done,转写模型的官方 usage 附在 completed 事件上,与 response.done 同等计费; // whisper 系按时长返回(token 全 0)时不走此分支,仍用本地估算兜底 - usage.TotalTokens += realtimeEvent.Usage.TotalTokens - usage.InputTokens += realtimeEvent.Usage.InputTokens - usage.OutputTokens += realtimeEvent.Usage.OutputTokens - usage.InputTokenDetails.AudioTokens += realtimeEvent.Usage.InputTokenDetails.AudioTokens - usage.InputTokenDetails.CachedTokens += realtimeEvent.Usage.InputTokenDetails.CachedTokens - usage.InputTokenDetails.TextTokens += realtimeEvent.Usage.InputTokenDetails.TextTokens - usage.OutputTokenDetails.AudioTokens += realtimeEvent.Usage.OutputTokenDetails.AudioTokens - // GA 转写 usage 不带 output_token_details,而计费公式只认明细字段,明细缺失时输出全按文本补记 - outTextTokens := realtimeEvent.Usage.OutputTokenDetails.TextTokens - if outTextTokens == 0 && realtimeEvent.Usage.OutputTokenDetails.AudioTokens == 0 { - outTextTokens = realtimeEvent.Usage.OutputTokens - } - usage.OutputTokenDetails.TextTokens += outTextTokens - if err := preConsumeUsage(c, info, usage, sumUsage); err != nil { + quotaResult, err := preConsumeUsage(c, info, transcriptionBilling.ModelName, &transcriptionBilling.Usage, transcriptionSumUsage) + // 已识别官方 usage 后清掉本句本地估算以防双计;这可能同时丢掉下一句已 append 的音频估算 + localUsage = &dto.RealtimeUsage{} + if err != nil { errChan <- fmt.Errorf("error consume usage: %v", err) return } - // 官方 usage 已入账,清掉本句攒下的本地估算,避免重复 - usage = &dto.RealtimeUsage{} - localUsage = &dto.RealtimeUsage{} + info.AddRealtimeTranscriptionQuota(quotaResult.Quota) + service.RecordRealtimeTranscriptionConsumeLog(c, info, transcriptionBilling.ModelName, &transcriptionBilling.Usage, quotaResult) } else if realtimeEvent.Type == dto.RealtimeEventTypeSessionUpdated || realtimeEvent.Type == dto.RealtimeEventTypeSessionCreated { realtimeSession := realtimeEvent.Session if realtimeSession != nil { // update audio format info.InputAudioFormat = common.GetStringIfEmpty(realtimeSession.InputAudioFormat, info.InputAudioFormat) info.OutputAudioFormat = common.GetStringIfEmpty(realtimeSession.OutputAudioFormat, info.OutputAudioFormat) + captureRealtimeTranscriptionModel(info, realtimeSession) } } else { textToken, audioToken, err := service.CountTokenRealtime(info, *realtimeEvent, info.UpstreamModelName) @@ -234,11 +220,11 @@ func OpenaiRealtimeHandler(c *gin.Context, info *relaycommon.RelayInfo) (*types. } if usage.TotalTokens != 0 { - _ = preConsumeUsage(c, info, usage, sumUsage) + _, _ = preConsumeUsage(c, info, info.OriginModelName, usage, sumUsage) } if localUsage.TotalTokens != 0 { - _ = preConsumeUsage(c, info, localUsage, sumUsage) + _, _ = preConsumeUsage(c, info, info.OriginModelName, localUsage, sumUsage) } // check usage total tokens, if 0, use local usage @@ -246,20 +232,62 @@ func OpenaiRealtimeHandler(c *gin.Context, info *relaycommon.RelayInfo) (*types. return nil, sumUsage } -func preConsumeUsage(ctx *gin.Context, info *relaycommon.RelayInfo, usage *dto.RealtimeUsage, totalUsage *dto.RealtimeUsage) error { +func addRealtimeUsage(total, delta dto.RealtimeUsage) dto.RealtimeUsage { + total.TotalTokens += delta.TotalTokens + total.InputTokens += delta.InputTokens + total.OutputTokens += delta.OutputTokens + total.InputTokenDetails.CachedTokens += delta.InputTokenDetails.CachedTokens + total.InputTokenDetails.TextTokens += delta.InputTokenDetails.TextTokens + total.InputTokenDetails.AudioTokens += delta.InputTokenDetails.AudioTokens + total.OutputTokenDetails.TextTokens += delta.OutputTokenDetails.TextTokens + total.OutputTokenDetails.AudioTokens += delta.OutputTokenDetails.AudioTokens + return total +} + +func billableRealtimeTranscriptionUsage(event *dto.RealtimeEvent) (dto.RealtimeUsage, bool) { + if event == nil || event.Type != dto.RealtimeEventInputAudioTranscriptionCompleted || event.Usage == nil || event.Usage.TotalTokens <= 0 { + return dto.RealtimeUsage{}, false + } + + usage := addRealtimeUsage(dto.RealtimeUsage{}, *event.Usage) + // GA 转写 usage 不带 output_token_details,而计费公式只认明细字段,明细缺失时输出全按文本补记 + if event.Usage.OutputTokenDetails.TextTokens == 0 && event.Usage.OutputTokenDetails.AudioTokens == 0 { + usage.OutputTokenDetails.TextTokens = event.Usage.OutputTokens + } + return usage, true +} + +type realtimeTranscriptionBillingInfo struct { + ModelName string + Usage dto.RealtimeUsage +} + +func realtimeTranscriptionBilling(originModelName, transcriptionModelName string, event *dto.RealtimeEvent) (realtimeTranscriptionBillingInfo, bool) { + usage, ok := billableRealtimeTranscriptionUsage(event) + if !ok { + return realtimeTranscriptionBillingInfo{}, false + } + if transcriptionModelName == "" { + transcriptionModelName = originModelName + } + return realtimeTranscriptionBillingInfo{ + ModelName: transcriptionModelName, + Usage: usage, + }, true +} + +func captureRealtimeTranscriptionModel(info *relaycommon.RelayInfo, session *dto.RealtimeSession) { + if info == nil || session == nil || session.InputAudioTranscription.Model == "" { + return + } + info.SetRealtimeTranscriptionModel(session.InputAudioTranscription.Model) +} + +func preConsumeUsage(ctx *gin.Context, info *relaycommon.RelayInfo, modelName string, usage *dto.RealtimeUsage, totalUsage *dto.RealtimeUsage) (service.WssQuotaResult, error) { if usage == nil || totalUsage == nil { - return fmt.Errorf("invalid usage pointer") + return service.WssQuotaResult{}, fmt.Errorf("invalid usage pointer") } - totalUsage.TotalTokens += usage.TotalTokens - totalUsage.InputTokens += usage.InputTokens - totalUsage.OutputTokens += usage.OutputTokens - totalUsage.InputTokenDetails.CachedTokens += usage.InputTokenDetails.CachedTokens - totalUsage.InputTokenDetails.TextTokens += usage.InputTokenDetails.TextTokens - totalUsage.InputTokenDetails.AudioTokens += usage.InputTokenDetails.AudioTokens - totalUsage.OutputTokenDetails.TextTokens += usage.OutputTokenDetails.TextTokens - totalUsage.OutputTokenDetails.AudioTokens += usage.OutputTokenDetails.AudioTokens - // clear usage - err := service.PreWssConsumeQuota(ctx, info, usage) - return err + *totalUsage = addRealtimeUsage(*totalUsage, *usage) + return service.PreWssConsumeQuota(ctx, info, modelName, usage) } diff --git a/relay/channel/openai/relay_realtime_test.go b/relay/channel/openai/relay_realtime_test.go new file mode 100644 index 000000000000..5cdb5896b8cf --- /dev/null +++ b/relay/channel/openai/relay_realtime_test.go @@ -0,0 +1,297 @@ +package openai + +import ( + "testing" + + "github.com/QuantumNous/new-api/common" + "github.com/QuantumNous/new-api/dto" + relaycommon "github.com/QuantumNous/new-api/relay/common" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func TestAddRealtimeUsage(t *testing.T) { + tests := []struct { + name string + total dto.RealtimeUsage + delta dto.RealtimeUsage + expected dto.RealtimeUsage + }{ + { + name: "adds every realtime billing field", + total: dto.RealtimeUsage{ + TotalTokens: 9, + InputTokens: 5, + OutputTokens: 4, + InputTokenDetails: dto.InputTokenDetails{ + CachedTokens: 1, + TextTokens: 2, + AudioTokens: 3, + }, + OutputTokenDetails: dto.OutputTokenDetails{TextTokens: 2, AudioTokens: 2}, + }, + delta: dto.RealtimeUsage{ + TotalTokens: 11, + InputTokens: 7, + OutputTokens: 4, + InputTokenDetails: dto.InputTokenDetails{ + CachedTokens: 2, + TextTokens: 3, + AudioTokens: 4, + }, + OutputTokenDetails: dto.OutputTokenDetails{TextTokens: 1, AudioTokens: 3}, + }, + expected: dto.RealtimeUsage{ + TotalTokens: 20, + InputTokens: 12, + OutputTokens: 8, + InputTokenDetails: dto.InputTokenDetails{ + CachedTokens: 3, + TextTokens: 5, + AudioTokens: 7, + }, + OutputTokenDetails: dto.OutputTokenDetails{TextTokens: 3, AudioTokens: 5}, + }, + }, + { + name: "copies billing fields into an empty total", + delta: dto.RealtimeUsage{ + TotalTokens: 7, + InputTokens: 6, + OutputTokens: 1, + InputTokenDetails: dto.InputTokenDetails{AudioTokens: 6}, + OutputTokenDetails: dto.OutputTokenDetails{ + TextTokens: 1, + }, + }, + expected: dto.RealtimeUsage{ + TotalTokens: 7, + InputTokens: 6, + OutputTokens: 1, + InputTokenDetails: dto.InputTokenDetails{AudioTokens: 6}, + OutputTokenDetails: dto.OutputTokenDetails{ + TextTokens: 1, + }, + }, + }, + } + + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + assert.Equal(t, test.expected, addRealtimeUsage(test.total, test.delta)) + }) + } +} + +func TestBillableRealtimeTranscriptionUsage(t *testing.T) { + tests := []struct { + name string + event *dto.RealtimeEvent + expected dto.RealtimeUsage + expectedBillable bool + }{ + { + name: "ignores a completed event without usage", + event: &dto.RealtimeEvent{ + Type: dto.RealtimeEventInputAudioTranscriptionCompleted, + }, + }, + { + name: "leaves whisper duration usage to local estimation", + event: &dto.RealtimeEvent{ + Type: dto.RealtimeEventInputAudioTranscriptionCompleted, + Usage: &dto.RealtimeUsage{}, + }, + }, + { + name: "ignores usage on a different event", + event: &dto.RealtimeEvent{ + Type: dto.RealtimeEventTypeResponseDone, + Usage: &dto.RealtimeUsage{TotalTokens: 1}, + }, + }, + { + name: "fills missing output details as text tokens", + event: &dto.RealtimeEvent{ + Type: dto.RealtimeEventInputAudioTranscriptionCompleted, + Usage: &dto.RealtimeUsage{ + TotalTokens: 13, + InputTokens: 10, + OutputTokens: 3, + InputTokenDetails: dto.InputTokenDetails{AudioTokens: 10}, + }, + }, + expected: dto.RealtimeUsage{ + TotalTokens: 13, + InputTokens: 10, + OutputTokens: 3, + InputTokenDetails: dto.InputTokenDetails{AudioTokens: 10}, + OutputTokenDetails: dto.OutputTokenDetails{ + TextTokens: 3, + }, + }, + expectedBillable: true, + }, + { + name: "preserves provided output details", + event: &dto.RealtimeEvent{ + Type: dto.RealtimeEventInputAudioTranscriptionCompleted, + Usage: &dto.RealtimeUsage{ + TotalTokens: 13, + InputTokens: 10, + OutputTokens: 3, + OutputTokenDetails: dto.OutputTokenDetails{ + TextTokens: 1, + AudioTokens: 2, + }, + }, + }, + expected: dto.RealtimeUsage{ + TotalTokens: 13, + InputTokens: 10, + OutputTokens: 3, + OutputTokenDetails: dto.OutputTokenDetails{ + TextTokens: 1, + AudioTokens: 2, + }, + }, + expectedBillable: true, + }, + { + name: "does not fill text when audio output details are present", + event: &dto.RealtimeEvent{ + Type: dto.RealtimeEventInputAudioTranscriptionCompleted, + Usage: &dto.RealtimeUsage{ + TotalTokens: 13, + InputTokens: 10, + OutputTokens: 3, + OutputTokenDetails: dto.OutputTokenDetails{ + AudioTokens: 3, + }, + }, + }, + expected: dto.RealtimeUsage{ + TotalTokens: 13, + InputTokens: 10, + OutputTokens: 3, + OutputTokenDetails: dto.OutputTokenDetails{ + AudioTokens: 3, + }, + }, + expectedBillable: true, + }, + } + + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + var original *dto.RealtimeUsage + if test.event != nil && test.event.Usage != nil { + usageCopy := *test.event.Usage + original = &usageCopy + } + + actual, billable := billableRealtimeTranscriptionUsage(test.event) + + assert.Equal(t, test.expectedBillable, billable) + assert.Equal(t, test.expected, actual) + if original != nil { + assert.Equal(t, *original, *test.event.Usage, "normalization must not mutate the upstream usage") + } + }) + } +} + +func TestCaptureRealtimeTranscriptionModel(t *testing.T) { + tests := []struct { + name string + payload string + currentModel string + expected string + }{ + { + name: "captures client session update", + payload: `{"type":"session.update","session":{"input_audio_transcription":{"model":"gpt-4o-transcribe"}}}`, + expected: "gpt-4o-transcribe", + }, + { + name: "captures upstream session created", + payload: `{"type":"session.created","session":{"input_audio_transcription":{"model":"gpt-4o-transcribe"}}}`, + expected: "gpt-4o-transcribe", + }, + { + name: "updates from upstream session updated", + payload: `{"type":"session.updated","session":{"input_audio_transcription":{"model":"gpt-4o-mini-transcribe"}}}`, + currentModel: "whisper-1", + expected: "gpt-4o-mini-transcribe", + }, + { + name: "does not overwrite with an empty model", + payload: `{"type":"session.updated","session":{"input_audio_transcription":{}}}`, + currentModel: "gpt-4o-transcribe", + expected: "gpt-4o-transcribe", + }, + { + name: "ignores a missing session", + payload: `{"type":"session.updated"}`, + currentModel: "gpt-4o-transcribe", + expected: "gpt-4o-transcribe", + }, + } + + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + var event dto.RealtimeEvent + require.NoError(t, common.Unmarshal([]byte(test.payload), &event)) + info := &relaycommon.RelayInfo{} + info.InitRealtimeTranscriptionState() + info.SetRealtimeTranscriptionModel(test.currentModel) + + captureRealtimeTranscriptionModel(info, event.Session) + + assert.Equal(t, test.expected, info.GetRealtimeTranscriptionModel()) + }) + } +} + +func TestRealtimeTranscriptionBilling(t *testing.T) { + event := &dto.RealtimeEvent{ + Type: dto.RealtimeEventInputAudioTranscriptionCompleted, + Usage: &dto.RealtimeUsage{ + TotalTokens: 13, + InputTokens: 10, + OutputTokens: 3, + InputTokenDetails: dto.InputTokenDetails{AudioTokens: 10}, + }, + } + tests := []struct { + name string + originModelName string + transcriptionModelName string + expectedModelName string + expectedOutputTextTokens int + }{ + { + name: "mixed session uses the ASR model and separate usage sum", + originModelName: "gpt-realtime", + transcriptionModelName: "gpt-4o-transcribe", + expectedModelName: "gpt-4o-transcribe", + expectedOutputTextTokens: 3, + }, + { + name: "transcription-only session falls back to the origin model", + originModelName: "gpt-4o-transcribe", + expectedModelName: "gpt-4o-transcribe", + expectedOutputTextTokens: 3, + }, + } + + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + billing, ok := realtimeTranscriptionBilling(test.originModelName, test.transcriptionModelName, event) + require.True(t, ok) + assert.Equal(t, test.expectedModelName, billing.ModelName) + assert.Equal(t, test.expectedOutputTextTokens, billing.Usage.OutputTokenDetails.TextTokens) + }) + } +} diff --git a/relay/common/realtime_transcription.go b/relay/common/realtime_transcription.go new file mode 100644 index 000000000000..591dc6c303f6 --- /dev/null +++ b/relay/common/realtime_transcription.go @@ -0,0 +1,67 @@ +package common + +import ( + "sync" + + basecommon "github.com/QuantumNous/new-api/common" + + "github.com/shopspring/decimal" +) + +// RealtimeTranscriptionState keeps ASR billing state synchronized between the +// client reader, upstream reader, and final websocket settlement. +type RealtimeTranscriptionState struct { + mu sync.RWMutex + model string + hasUsage bool + quota decimal.Decimal +} + +func (info *RelayInfo) InitRealtimeTranscriptionState() { + if info == nil { + return + } + if info.RealtimeTranscription == nil { + info.RealtimeTranscription = &RealtimeTranscriptionState{} + } +} + +func (info *RelayInfo) SetRealtimeTranscriptionModel(model string) { + if info == nil || model == "" { + return + } + info.InitRealtimeTranscriptionState() + info.RealtimeTranscription.mu.Lock() + info.RealtimeTranscription.model = model + info.RealtimeTranscription.mu.Unlock() +} + +func (info *RelayInfo) GetRealtimeTranscriptionModel() string { + if info == nil || info.RealtimeTranscription == nil { + return "" + } + info.RealtimeTranscription.mu.RLock() + defer info.RealtimeTranscription.mu.RUnlock() + return info.RealtimeTranscription.model +} + +func (info *RelayInfo) AddRealtimeTranscriptionQuota(quota int) { + if info == nil { + return + } + info.InitRealtimeTranscriptionState() + info.RealtimeTranscription.mu.Lock() + info.RealtimeTranscription.hasUsage = true + info.RealtimeTranscription.quota = info.RealtimeTranscription.quota.Add(decimal.NewFromInt(int64(quota))) + info.RealtimeTranscription.mu.Unlock() +} + +func (info *RelayInfo) GetRealtimeTranscriptionBilling() (bool, int, *basecommon.QuotaClamp) { + if info == nil || info.RealtimeTranscription == nil { + return false, 0, nil + } + info.RealtimeTranscription.mu.RLock() + defer info.RealtimeTranscription.mu.RUnlock() + quota, clamp := basecommon.QuotaFromDecimalChecked(info.RealtimeTranscription.quota) + return info.RealtimeTranscription.hasUsage, quota, clamp +} diff --git a/relay/common/realtime_transcription_test.go b/relay/common/realtime_transcription_test.go new file mode 100644 index 000000000000..584cd8b4a990 --- /dev/null +++ b/relay/common/realtime_transcription_test.go @@ -0,0 +1,40 @@ +package common + +import ( + "testing" + + basecommon "github.com/QuantumNous/new-api/common" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func TestRealtimeTranscriptionState(t *testing.T) { + info := &RelayInfo{} + info.InitRealtimeTranscriptionState() + info.SetRealtimeTranscriptionModel("gpt-4o-transcribe") + info.SetRealtimeTranscriptionModel("") + info.AddRealtimeTranscriptionQuota(10) + info.AddRealtimeTranscriptionQuota(20) + + hasUsage, quota, clamp := info.GetRealtimeTranscriptionBilling() + + require.True(t, hasUsage) + assert.Equal(t, "gpt-4o-transcribe", info.GetRealtimeTranscriptionModel()) + assert.Equal(t, 30, quota) + assert.Nil(t, clamp) +} + +func TestRealtimeTranscriptionQuotaSaturatesForFinalStatistics(t *testing.T) { + info := &RelayInfo{} + info.InitRealtimeTranscriptionState() + info.AddRealtimeTranscriptionQuota(basecommon.MaxQuota) + info.AddRealtimeTranscriptionQuota(1) + + hasUsage, quota, clamp := info.GetRealtimeTranscriptionBilling() + + require.True(t, hasUsage) + assert.Equal(t, basecommon.MaxQuota, quota) + require.NotNil(t, clamp) + assert.Equal(t, basecommon.QuotaClampOverflow, clamp.Kind) +} diff --git a/relay/common/relay_info.go b/relay/common/relay_info.go index 9f460ce5c6a7..1bb646ffd16b 100644 --- a/relay/common/relay_info.go +++ b/relay/common/relay_info.go @@ -122,6 +122,10 @@ type RelayInfo struct { SendResponseCount int ReceivedResponseCount int FinalPreConsumedQuota int // 最终预消耗的配额 + + // Realtime 转写 usage 独立增量扣费,不进入主模型的最终结算 + RealtimeTranscription *RealtimeTranscriptionState + // ForcePreConsume 为 true 时禁用 BillingSession 的信任额度旁路， // 强制预扣全额。用于异步任务（视频/音乐生成等），因为请求返回后任务仍在运行， // 必须在提交前锁定全额。 diff --git a/service/quota.go b/service/quota.go index e5d4ec7e2330..7ca9356cb05f 100644 --- a/service/quota.go +++ b/service/quota.go @@ -39,6 +39,33 @@ type QuotaInfo struct { GroupRatio float64 } +// WssQuotaResult keeps an incremental websocket charge and its ratio snapshot +// together, so consume logging does not perform a second settings lookup. +type WssQuotaResult struct { + Quota int + ModelRatio float64 + GroupRatio float64 + UserGroupRatio float64 + CompletionRatio float64 + AudioRatio float64 + AudioCompletionRatio float64 +} + +type wssQuotaSummary struct { + SettlementQuota int + StatisticsQuota int +} + +func summarizeWssQuotas(mainQuota, transcriptionQuota int) (wssQuotaSummary, *common.QuotaClamp) { + statisticsQuota, clamp := common.QuotaFromDecimalChecked( + decimal.NewFromInt(int64(mainQuota)).Add(decimal.NewFromInt(int64(transcriptionQuota))), + ) + return wssQuotaSummary{ + SettlementQuota: mainQuota, + StatisticsQuota: statisticsQuota, + }, clamp +} + func hasCustomModelRatio(modelName string, currentRatio float64) bool { defaultRatio, exists := ratio_setting.GetDefaultModelRatioMap()[modelName] if !exists { @@ -86,40 +113,45 @@ func calculateAudioQuota(info QuotaInfo) (int, *common.QuotaClamp) { return common.QuotaFromDecimalChecked(quota) } -func PreWssConsumeQuota(ctx *gin.Context, relayInfo *relaycommon.RelayInfo, usage *dto.RealtimeUsage) error { +func getWssGroupRatio(ctx *gin.Context, relayInfo *relaycommon.RelayInfo) (float64, float64) { + autoGroup, exists := common.GetContextKey(ctx, constant.ContextKeyAutoGroup) + if exists { + relayInfo.UsingGroup = autoGroup.(string) + logger.LogDebug(ctx, "final group ratio: %f", ratio_setting.GetGroupRatio(relayInfo.UsingGroup)) + } + + groupRatio := ratio_setting.GetGroupRatio(relayInfo.UsingGroup) + userGroupRatio := -1.0 + if ratio, ok := ratio_setting.GetGroupGroupRatio(relayInfo.UserGroup, relayInfo.UsingGroup); ok { + groupRatio = ratio + userGroupRatio = ratio + } + return groupRatio, userGroupRatio +} + +func PreWssConsumeQuota(ctx *gin.Context, relayInfo *relaycommon.RelayInfo, modelName string, usage *dto.RealtimeUsage) (WssQuotaResult, error) { if relayInfo.UsePrice { - return nil + return WssQuotaResult{}, nil } userQuota, err := model.GetUserQuota(relayInfo.UserId, false) if err != nil { - return err + return WssQuotaResult{}, err } token, err := model.GetTokenByKey(strings.TrimPrefix(relayInfo.TokenKey, "sk-"), false) if err != nil { - return err + return WssQuotaResult{}, err } - modelName := relayInfo.OriginModelName textInputTokens := usage.InputTokenDetails.TextTokens textOutTokens := usage.OutputTokenDetails.TextTokens audioInputTokens := usage.InputTokenDetails.AudioTokens audioOutTokens := usage.OutputTokenDetails.AudioTokens - groupRatio := ratio_setting.GetGroupRatio(relayInfo.UsingGroup) modelRatio, _, _ := ratio_setting.GetModelRatio(modelName) - - autoGroup, exists := common.GetContextKey(ctx, constant.ContextKeyAutoGroup) - if exists { - groupRatio = ratio_setting.GetGroupRatio(autoGroup.(string)) - logger.LogDebug(ctx, "final group ratio: %f", groupRatio) - relayInfo.UsingGroup = autoGroup.(string) - } - - actualGroupRatio := groupRatio - userGroupRatio, ok := ratio_setting.GetGroupGroupRatio(relayInfo.UserGroup, relayInfo.UsingGroup) - if ok { - actualGroupRatio = userGroupRatio - } + actualGroupRatio, userGroupRatio := getWssGroupRatio(ctx, relayInfo) + completionRatio := ratio_setting.GetCompletionRatio(modelName) + audioRatio := ratio_setting.GetAudioRatio(modelName) + audioCompletionRatio := ratio_setting.GetAudioCompletionRatio(modelName) quotaInfo := QuotaInfo{ InputDetails: TokenDetails{ @@ -138,21 +170,68 @@ func PreWssConsumeQuota(ctx *gin.Context, relayInfo *relaycommon.RelayInfo, usag quota, clamp := calculateAudioQuota(quotaInfo) noteQuotaClamp(relayInfo, clamp) + result := WssQuotaResult{ + Quota: quota, + ModelRatio: modelRatio, + GroupRatio: actualGroupRatio, + UserGroupRatio: userGroupRatio, + CompletionRatio: completionRatio, + AudioRatio: audioRatio, + AudioCompletionRatio: audioCompletionRatio, + } if userQuota < quota { - return fmt.Errorf("user quota is not enough, user quota: %s, need quota: %s", logger.FormatQuota(userQuota), logger.FormatQuota(quota)) + return WssQuotaResult{}, fmt.Errorf("user quota is not enough, user quota: %s, need quota: %s", logger.FormatQuota(userQuota), logger.FormatQuota(quota)) } if !token.UnlimitedQuota && token.RemainQuota < quota { - return fmt.Errorf("token quota is not enough, token remain quota: %s, need quota: %s", logger.FormatQuota(token.RemainQuota), logger.FormatQuota(quota)) + return WssQuotaResult{}, fmt.Errorf("token quota is not enough, token remain quota: %s, need quota: %s", logger.FormatQuota(token.RemainQuota), logger.FormatQuota(quota)) } err = PostConsumeQuota(relayInfo, quota, 0, false) if err != nil { - return err + return WssQuotaResult{}, err } logger.LogInfo(ctx, "realtime streaming consume quota success, quota: "+fmt.Sprintf("%d", quota)) - return nil + return result, nil +} + +// RecordRealtimeTranscriptionConsumeLog records an ASR charge already deducted +// by PreWssConsumeQuota. It must not settle BillingSession or update aggregate +// request statistics; PostWssConsumeQuota performs those operations once per session. +func RecordRealtimeTranscriptionConsumeLog(ctx *gin.Context, relayInfo *relaycommon.RelayInfo, modelName string, + usage *dto.RealtimeUsage, result WssQuotaResult) { + if relayInfo == nil || usage == nil { + return + } + + logContent := fmt.Sprintf("模型倍率 %.2f，补全倍率 %.2f，音频倍率 %.2f，音频补全倍率 %.2f，分组倍率 %.2f", + result.ModelRatio, result.CompletionRatio, result.AudioRatio, result.AudioCompletionRatio, result.GroupRatio) + logContent += ", 实时转写" + + other := GenerateWssOtherInfo(ctx, relayInfo, usage, result.ModelRatio, result.GroupRatio, + result.CompletionRatio, result.AudioRatio, result.AudioCompletionRatio, 0, result.UserGroupRatio) + // 增量日志早于主 BillingSession 最终结算,不能记录尚未完成的会话级订阅汇总 + delete(other, "subscription_pre_consumed") + delete(other, "subscription_post_delta") + delete(other, "subscription_consumed") + delete(other, "subscription_used") + delete(other, "subscription_remain") + attachQuotaSaturation(ctx, relayInfo, other) + model.RecordConsumeLog(ctx, relayInfo.UserId, model.RecordConsumeLogParams{ + ChannelId: relayInfo.ChannelId, + PromptTokens: usage.InputTokens, + CompletionTokens: usage.OutputTokens, + ModelName: modelName, + TokenName: ctx.GetString("token_name"), + Quota: result.Quota, + Content: logContent, + TokenId: relayInfo.TokenId, + UseTimeSeconds: int(time.Now().Unix() - relayInfo.StartTime.Unix()), + IsStream: relayInfo.IsStream, + Group: relayInfo.UsingGroup, + Other: other, + }) } func PostWssConsumeQuota(ctx *gin.Context, relayInfo *relaycommon.RelayInfo, modelName string, @@ -207,6 +286,9 @@ func PostWssConsumeQuota(ctx *gin.Context, relayInfo *relaycommon.RelayInfo, mod } totalTokens := usage.TotalTokens + hasMainUsage := totalTokens != 0 + hasTranscriptionUsage, transcriptionQuota, transcriptionClamp := relayInfo.GetRealtimeTranscriptionBilling() + noteQuotaClamp(relayInfo, transcriptionClamp) var logContent string if !usePrice { logContent = fmt.Sprintf("模型倍率 %.2f，补全倍率 %.2f，音频倍率 %.2f，音频补全倍率 %.2f，分组倍率 %.2f", @@ -216,21 +298,28 @@ func PostWssConsumeQuota(ctx *gin.Context, relayInfo *relaycommon.RelayInfo, mod } // record all the consume log even if quota is 0 - if totalTokens == 0 { - // in this case, must be some error happened - // we cannot just return, because we may have to return the pre-consumed quota + if !hasMainUsage { + // 主 usage 为空时仍需用 0 结算,退回可能存在的主模型预扣费 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, modelName, relayInfo.FinalPreConsumedQuota)) - } else { - model.UpdateUserUsedQuotaAndRequestCount(relayInfo.UserId, quota) - model.UpdateChannelUsedQuota(relayInfo.ChannelId, quota) + if !hasTranscriptionUsage { + 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, modelName, relayInfo.FinalPreConsumedQuota)) + } + } + quotaSummary, summaryClamp := summarizeWssQuotas(quota, transcriptionQuota) + noteQuotaClamp(relayInfo, summaryClamp) + if hasMainUsage || hasTranscriptionUsage { + model.UpdateUserUsedQuotaAndRequestCount(relayInfo.UserId, quotaSummary.StatisticsQuota) + model.UpdateChannelUsedQuota(relayInfo.ChannelId, quotaSummary.StatisticsQuota) } - if err := SettleBilling(ctx, relayInfo, quota); err != nil { + if err := SettleBilling(ctx, relayInfo, quotaSummary.SettlementQuota); err != nil { logger.LogError(ctx, "error settling billing: "+err.Error()) } + if !hasMainUsage && hasTranscriptionUsage { + return + } logModel := modelName if extraContent != "" { diff --git a/service/quota_realtime_test.go b/service/quota_realtime_test.go new file mode 100644 index 000000000000..6832d5dd8b6f --- /dev/null +++ b/service/quota_realtime_test.go @@ -0,0 +1,63 @@ +package service + +import ( + "testing" + + "github.com/QuantumNous/new-api/common" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func TestSummarizeWssQuotasKeepsTranscriptionOutOfSettlement(t *testing.T) { + tests := []struct { + name string + mainQuota int + transcriptionQuota int + expected wssQuotaSummary + }{ + { + name: "mixed realtime and transcription session", + mainQuota: 100, + transcriptionQuota: 25, + expected: wssQuotaSummary{ + SettlementQuota: 100, + StatisticsQuota: 125, + }, + }, + { + name: "transcription-only session", + transcriptionQuota: 25, + expected: wssQuotaSummary{ + SettlementQuota: 0, + StatisticsQuota: 25, + }, + }, + { + name: "main realtime session without transcription", + mainQuota: 100, + expected: wssQuotaSummary{ + SettlementQuota: 100, + StatisticsQuota: 100, + }, + }, + } + + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + actual, clamp := summarizeWssQuotas(test.mainQuota, test.transcriptionQuota) + + assert.Equal(t, test.expected, actual) + assert.Nil(t, clamp) + }) + } +} + +func TestSummarizeWssQuotasSaturatesStatisticsOnly(t *testing.T) { + summary, clamp := summarizeWssQuotas(common.MaxQuota, 1) + + assert.Equal(t, common.MaxQuota, summary.SettlementQuota) + assert.Equal(t, common.MaxQuota, summary.StatisticsQuota) + require.NotNil(t, clamp) + assert.Equal(t, common.QuotaClampOverflow, clamp.Kind) +}