Skip to content

Commit e052d7d

Browse files
Calcium-IonCalcium-Ion
committed
feat: realtime pre consume
(cherry picked from commit d87917f8f6eb9d2e144a9f840d6d91767ea2eb69)
1 parent efd7ea0 commit e052d7d

4 files changed

Lines changed: 180 additions & 94 deletions

File tree

‎relay/channel/openai/relay-openai.go‎

Lines changed: 45 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -389,6 +389,7 @@ func OpenaiRealtimeHandler(c *gin.Context, info *relaycommon.RelayInfo) (*dto.Op
389389

390390
usage := &dto.RealtimeUsage{}
391391
localUsage := &dto.RealtimeUsage{}
392+
sumUsage := &dto.RealtimeUsage{}
392393

393394
go func() {
394395
for {
@@ -478,6 +479,12 @@ func OpenaiRealtimeHandler(c *gin.Context, info *relaycommon.RelayInfo) (*dto.Op
478479
usage.InputTokenDetails.TextTokens += realtimeUsage.InputTokenDetails.TextTokens
479480
usage.OutputTokenDetails.AudioTokens += realtimeUsage.OutputTokenDetails.AudioTokens
480481
usage.OutputTokenDetails.TextTokens += realtimeUsage.OutputTokenDetails.TextTokens
482+
err := preConsumeUsage(c, info, usage, sumUsage)
483+
if err != nil {
484+
errChan <- fmt.Errorf("error consume usage: %v", err)
485+
return
486+
}
487+
usage = &dto.RealtimeUsage{}
481488
} else {
482489
textToken, audioToken, err := service.CountTokenRealtime(info, *realtimeEvent, info.UpstreamModelName)
483490
if err != nil {
@@ -490,7 +497,18 @@ func OpenaiRealtimeHandler(c *gin.Context, info *relaycommon.RelayInfo) (*dto.Op
490497
localUsage.InputTokens += textToken + audioToken
491498
localUsage.InputTokenDetails.TextTokens += textToken
492499
localUsage.InputTokenDetails.AudioTokens += audioToken
500+
err = preConsumeUsage(c, info, localUsage, sumUsage)
501+
if err != nil {
502+
errChan <- fmt.Errorf("error consume usage: %v", err)
503+
return
504+
}
505+
localUsage = &dto.RealtimeUsage{}
506+
// print now usage
493507
}
508+
common.LogInfo(c, fmt.Sprintf("realtime streaming sumUsage: %v", sumUsage))
509+
common.LogInfo(c, fmt.Sprintf("realtime streaming localUsage: %v", localUsage))
510+
common.LogInfo(c, fmt.Sprintf("realtime streaming localUsage: %v", localUsage))
511+
494512
} else if realtimeEvent.Type == dto.RealtimeEventTypeSessionUpdated || realtimeEvent.Type == dto.RealtimeEventTypeSessionCreated {
495513
realtimeSession := realtimeEvent.Session
496514
if realtimeSession != nil {
@@ -528,15 +546,38 @@ func OpenaiRealtimeHandler(c *gin.Context, info *relaycommon.RelayInfo) (*dto.Op
528546
select {
529547
case <-clientClosed:
530548
case <-targetClosed:
531-
case <-errChan:
549+
case err := <-errChan:
532550
//return service.OpenAIErrorWrapper(err, "realtime_error", http.StatusInternalServerError), nil
551+
common.LogError(c, "realtime error: "+err.Error())
533552
case <-c.Done():
534553
}
535554

555+
if usage.TotalTokens != 0 {
556+
_ = preConsumeUsage(c, info, usage, sumUsage)
557+
}
558+
559+
if localUsage.TotalTokens != 0 {
560+
_ = preConsumeUsage(c, info, localUsage, sumUsage)
561+
}
562+
536563
// check usage total tokens, if 0, use local usage
537564

538-
if usage.TotalTokens == 0 {
539-
usage = localUsage
565+
return nil, sumUsage
566+
}
567+
568+
func preConsumeUsage(ctx *gin.Context, info *relaycommon.RelayInfo, usage *dto.RealtimeUsage, totalUsage *dto.RealtimeUsage) error {
569+
totalUsage.TotalTokens += usage.TotalTokens
570+
totalUsage.InputTokens += usage.InputTokens
571+
totalUsage.OutputTokens += usage.OutputTokens
572+
totalUsage.InputTokenDetails.CachedTokens += usage.InputTokenDetails.CachedTokens
573+
totalUsage.InputTokenDetails.TextTokens += usage.InputTokenDetails.TextTokens
574+
totalUsage.InputTokenDetails.AudioTokens += usage.InputTokenDetails.AudioTokens
575+
totalUsage.OutputTokenDetails.TextTokens += usage.OutputTokenDetails.TextTokens
576+
totalUsage.OutputTokenDetails.AudioTokens += usage.OutputTokenDetails.AudioTokens
577+
// clear usage
578+
err := service.PreWssConsumeQuota(ctx, info, usage)
579+
if err == nil {
580+
common.LogInfo(ctx, "realtime streaming consume usage success")
540581
}
541-
return nil, usage
582+
return err
542583
}

‎relay/common/relay_info.go‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -23,6 +23,7 @@ type RelayInfo struct {
2323
ApiType int
2424
IsStream bool
2525
IsPlayground bool
26+
UsePrice bool
2627
RelayMode int
2728
UpstreamModelName string
2829
OriginModelName string

‎relay/websocket.go‎

Lines changed: 2 additions & 90 deletions
Original file line numberDiff line numberDiff line change
@@ -5,15 +5,11 @@ import (
55
"fmt"
66
"github.com/gin-gonic/gin"
77
"github.com/gorilla/websocket"
8-
"math"
98
"net/http"
109
"one-api/common"
1110
"one-api/dto"
12-
"one-api/model"
1311
relaycommon "one-api/relay/common"
1412
"one-api/service"
15-
"strings"
16-
"time"
1713
)
1814

1915
//func getAndValidateWssRequest(c *gin.Context, ws *websocket.Conn) (*dto.RealtimeEvent, error) {
@@ -91,6 +87,7 @@ func WssHelper(c *gin.Context, ws *websocket.Conn) *dto.OpenAIErrorWithStatusCod
9187
preConsumedQuota = int(float64(preConsumedTokens) * ratio)
9288
} else {
9389
preConsumedQuota = int(modelPrice * common.QuotaPerUnit * groupRatio)
90+
relayInfo.UsePrice = true
9491
}
9592

9693
// pre-consume quota 预消耗配额
@@ -126,95 +123,10 @@ func WssHelper(c *gin.Context, ws *websocket.Conn) *dto.OpenAIErrorWithStatusCod
126123
service.ResetStatusCode(openaiErr, statusCodeMappingStr)
127124
return openaiErr
128125
}
129-
postWssConsumeQuota(c, relayInfo, relayInfo.UpstreamModelName, usage.(*dto.RealtimeUsage), ratio, preConsumedQuota, userQuota, modelRatio, groupRatio, modelPrice, getModelPriceSuccess, "")
126+
service.PostWssConsumeQuota(c, relayInfo, relayInfo.UpstreamModelName, usage.(*dto.RealtimeUsage), ratio, preConsumedQuota, userQuota, modelRatio, groupRatio, modelPrice, getModelPriceSuccess, "")
130127
return nil
131128
}
132129

133-
func postWssConsumeQuota(ctx *gin.Context, relayInfo *relaycommon.RelayInfo, modelName string,
134-
usage *dto.RealtimeUsage, ratio float64, preConsumedQuota int, userQuota int, modelRatio float64,
135-
groupRatio float64,
136-
modelPrice float64, usePrice bool, extraContent string) {
137-
138-
useTimeSeconds := time.Now().Unix() - relayInfo.StartTime.Unix()
139-
textInputTokens := usage.InputTokenDetails.TextTokens
140-
textOutTokens := usage.OutputTokenDetails.TextTokens
141-
142-
audioInputTokens := usage.InputTokenDetails.AudioTokens
143-
audioOutTokens := usage.OutputTokenDetails.AudioTokens
144-
145-
tokenName := ctx.GetString("token_name")
146-
completionRatio := common.GetCompletionRatio(modelName)
147-
audioRatio := common.GetAudioRatio(relayInfo.UpstreamModelName)
148-
audioCompletionRatio := common.GetAudioCompletionRatio(modelName)
149-
150-
quota := 0
151-
if !usePrice {
152-
quota = textInputTokens + int(math.Round(float64(textOutTokens)*completionRatio))
153-
quota += int(math.Round(float64(audioInputTokens)*audioRatio)) + int(math.Round(float64(audioOutTokens)*audioRatio*audioCompletionRatio))
154-
155-
quota = int(math.Round(float64(quota) * ratio))
156-
if ratio != 0 && quota <= 0 {
157-
quota = 1
158-
}
159-
} else {
160-
quota = int(modelPrice * common.QuotaPerUnit * groupRatio)
161-
}
162-
totalTokens := usage.TotalTokens
163-
var logContent string
164-
if !usePrice {
165-
logContent = fmt.Sprintf("模型倍率 %.2f，补全倍率 %.2f，音频倍率 %.2f，音频补全倍率 %.2f，分组倍率 %.2f", modelRatio, completionRatio, audioRatio, audioCompletionRatio, groupRatio)
166-
} else {
167-
logContent = fmt.Sprintf("模型价格 %.2f，分组倍率 %.2f", modelPrice, groupRatio)
168-
}
169-
170-
// record all the consume log even if quota is 0
171-
if totalTokens == 0 {
172-
// in this case, must be some error happened
173-
// we cannot just return, because we may have to return the pre-consumed quota
174-
quota = 0
175-
logContent += fmt.Sprintf("（可能是上游超时）")
176-
common.LogError(ctx, fmt.Sprintf("total tokens is 0, cannot consume quota, userId %d, channelId %d, "+
177-
"tokenId %d, model %s， pre-consumed quota %d", relayInfo.UserId, relayInfo.ChannelId, relayInfo.TokenId, modelName, preConsumedQuota))
178-
} else {
179-
//if sensitiveResp != nil {
180-
// logContent += fmt.Sprintf("，敏感词：%s", strings.Join(sensitiveResp.SensitiveWords, ", "))
181-
//}
182-
quotaDelta := quota - preConsumedQuota
183-
if quotaDelta != 0 {
184-
err := model.PostConsumeTokenQuota(relayInfo, userQuota, quotaDelta, preConsumedQuota, true)
185-
if err != nil {
186-
common.LogError(ctx, "error consuming token remain quota: "+err.Error())
187-
}
188-
}
189-
err := model.CacheUpdateUserQuota(relayInfo.UserId)
190-
if err != nil {
191-
common.LogError(ctx, "error update user quota cache: "+err.Error())
192-
}
193-
model.UpdateUserUsedQuotaAndRequestCount(relayInfo.UserId, quota)
194-
model.UpdateChannelUsedQuota(relayInfo.ChannelId, quota)
195-
}
196-
197-
logModel := modelName
198-
if strings.HasPrefix(logModel, "gpt-4-gizmo") {
199-
logModel = "gpt-4-gizmo-*"
200-
logContent += fmt.Sprintf("，模型 %s", modelName)
201-
}
202-
if strings.HasPrefix(logModel, "gpt-4o-gizmo") {
203-
logModel = "gpt-4o-gizmo-*"
204-
logContent += fmt.Sprintf("，模型 %s", modelName)
205-
}
206-
if extraContent != "" {
207-
logContent += ", " + extraContent
208-
}
209-
other := service.GenerateWssOtherInfo(ctx, relayInfo, usage, modelRatio, groupRatio, completionRatio, modelPrice)
210-
model.RecordConsumeLog(ctx, relayInfo.UserId, relayInfo.ChannelId, usage.InputTokens, usage.OutputTokens, logModel,
211-
tokenName, quota, logContent, relayInfo.TokenId, userQuota, int(useTimeSeconds), relayInfo.IsStream, other)
212-
213-
//if quota != 0 {
214-
//
215-
//}
216-
}
217-
218130
//func getWssPromptTokens(textRequest *dto.RealtimeEvent, info *relaycommon.RelayInfo) (int, error) {
219131
// var promptTokens int
220132
// var err error

‎service/quota.go‎

Lines changed: 132 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,132 @@
1+
package service
2+
3+
import (
4+
"fmt"
5+
"github.com/gin-gonic/gin"
6+
"math"
7+
"one-api/common"
8+
"one-api/dto"
9+
"one-api/model"
10+
relaycommon "one-api/relay/common"
11+
"strings"
12+
"time"
13+
)
14+
15+
func PreWssConsumeQuota(ctx *gin.Context, relayInfo *relaycommon.RelayInfo, usage *dto.RealtimeUsage) error {
16+
if relayInfo.UsePrice {
17+
return nil
18+
}
19+
modelName := relayInfo.UpstreamModelName
20+
textInputTokens := usage.InputTokenDetails.TextTokens
21+
textOutTokens := usage.OutputTokenDetails.TextTokens
22+
audioInputTokens := usage.InputTokenDetails.AudioTokens
23+
audioOutTokens := usage.OutputTokenDetails.AudioTokens
24+
25+
completionRatio := common.GetCompletionRatio(modelName)
26+
audioRatio := common.GetAudioRatio(relayInfo.UpstreamModelName)
27+
audioCompletionRatio := common.GetAudioCompletionRatio(modelName)
28+
groupRatio := common.GetGroupRatio(relayInfo.Group)
29+
modelRatio := common.GetModelRatio(modelName)
30+
31+
ratio := groupRatio * modelRatio
32+
33+
quota := textInputTokens + int(math.Round(float64(textOutTokens)*completionRatio))
34+
quota += int(math.Round(float64(audioInputTokens)*audioRatio)) + int(math.Round(float64(audioOutTokens)*audioRatio*audioCompletionRatio))
35+
36+
quota = int(math.Round(float64(quota) * ratio))
37+
if ratio != 0 && quota <= 0 {
38+
quota = 1
39+
}
40+
41+
err := model.PostConsumeTokenQuota(relayInfo, 0, quota, 0, false)
42+
if err != nil {
43+
return err
44+
}
45+
err = model.CacheUpdateUserQuota(relayInfo.UserId)
46+
if err != nil {
47+
return err
48+
}
49+
return nil
50+
}
51+
52+
func PostWssConsumeQuota(ctx *gin.Context, relayInfo *relaycommon.RelayInfo, modelName string,
53+
usage *dto.RealtimeUsage, ratio float64, preConsumedQuota int, userQuota int, modelRatio float64,
54+
groupRatio float64,
55+
modelPrice float64, usePrice bool, extraContent string) {
56+
57+
useTimeSeconds := time.Now().Unix() - relayInfo.StartTime.Unix()
58+
textInputTokens := usage.InputTokenDetails.TextTokens
59+
textOutTokens := usage.OutputTokenDetails.TextTokens
60+
61+
audioInputTokens := usage.InputTokenDetails.AudioTokens
62+
audioOutTokens := usage.OutputTokenDetails.AudioTokens
63+
64+
tokenName := ctx.GetString("token_name")
65+
completionRatio := common.GetCompletionRatio(modelName)
66+
audioRatio := common.GetAudioRatio(relayInfo.UpstreamModelName)
67+
audioCompletionRatio := common.GetAudioCompletionRatio(modelName)
68+
69+
quota := 0
70+
if !usePrice {
71+
quota = textInputTokens + int(math.Round(float64(textOutTokens)*completionRatio))
72+
quota += int(math.Round(float64(audioInputTokens)*audioRatio)) + int(math.Round(float64(audioOutTokens)*audioRatio*audioCompletionRatio))
73+
74+
quota = int(math.Round(float64(quota) * ratio))
75+
if ratio != 0 && quota <= 0 {
76+
quota = 1
77+
}
78+
} else {
79+
quota = int(modelPrice * common.QuotaPerUnit * groupRatio)
80+
}
81+
totalTokens := usage.TotalTokens
82+
var logContent string
83+
if !usePrice {
84+
logContent = fmt.Sprintf("模型倍率 %.2f，补全倍率 %.2f，音频倍率 %.2f，音频补全倍率 %.2f，分组倍率 %.2f", modelRatio, completionRatio, audioRatio, audioCompletionRatio, groupRatio)
85+
} else {
86+
logContent = fmt.Sprintf("模型价格 %.2f，分组倍率 %.2f", modelPrice, groupRatio)
87+
}
88+
89+
// record all the consume log even if quota is 0
90+
if totalTokens == 0 {
91+
// in this case, must be some error happened
92+
// we cannot just return, because we may have to return the pre-consumed quota
93+
quota = 0
94+
logContent += fmt.Sprintf("（可能是上游超时）")
95+
common.LogError(ctx, fmt.Sprintf("total tokens is 0, cannot consume quota, userId %d, channelId %d, "+
96+
"tokenId %d, model %s， pre-consumed quota %d", relayInfo.UserId, relayInfo.ChannelId, relayInfo.TokenId, modelName, preConsumedQuota))
97+
} else {
98+
//if sensitiveResp != nil {
99+
// logContent += fmt.Sprintf("，敏感词：%s", strings.Join(sensitiveResp.SensitiveWords, ", "))
100+
//}
101+
//quotaDelta := quota - preConsumedQuota
102+
//if quotaDelta != 0 {
103+
// err := model.PostConsumeTokenQuota(relayInfo, userQuota, quotaDelta, preConsumedQuota, true)
104+
// if err != nil {
105+
// common.LogError(ctx, "error consuming token remain quota: "+err.Error())
106+
// }
107+
//}
108+
109+
//err := model.CacheUpdateUserQuota(relayInfo.UserId)
110+
//if err != nil {
111+
// common.LogError(ctx, "error update user quota cache: "+err.Error())
112+
//}
113+
model.UpdateUserUsedQuotaAndRequestCount(relayInfo.UserId, quota)
114+
model.UpdateChannelUsedQuota(relayInfo.ChannelId, quota)
115+
}
116+
117+
logModel := modelName
118+
if strings.HasPrefix(logModel, "gpt-4-gizmo") {
119+
logModel = "gpt-4-gizmo-*"
120+
logContent += fmt.Sprintf("，模型 %s", modelName)
121+
}
122+
if strings.HasPrefix(logModel, "gpt-4o-gizmo") {
123+
logModel = "gpt-4o-gizmo-*"
124+
logContent += fmt.Sprintf("，模型 %s", modelName)
125+
}
126+
if extraContent != "" {
127+
logContent += ", " + extraContent
128+
}
129+
other := GenerateWssOtherInfo(ctx, relayInfo, usage, modelRatio, groupRatio, completionRatio, modelPrice)
130+
model.RecordConsumeLog(ctx, relayInfo.UserId, relayInfo.ChannelId, usage.InputTokens, usage.OutputTokens, logModel,
131+
tokenName, quota, logContent, relayInfo.TokenId, userQuota, int(useTimeSeconds), relayInfo.IsStream, other)
132+
}

0 commit comments

Comments
 (0)