mirror of
https://github.com/Wei-Shaw/sub2api.git
synced 2026-09-24 16:05:44 +08:00
Merge pull request #3534 from mxyhi/fix/openai-quota-platform-post-billing
fix(openai): preserve quota platform in usage billing
This commit is contained in:
@@ -298,6 +298,7 @@ func (h *OpenAIGatewayHandler) ChatCompletions(c *gin.Context) {
|
||||
clientIP := ip.GetClientIP(c)
|
||||
inboundEndpoint := GetInboundEndpoint(c)
|
||||
upstreamEndpoint := resolveOpenAIUpstreamEndpoint(c, account)
|
||||
quotaPlatform := service.QuotaPlatform(c.Request.Context(), apiKey)
|
||||
|
||||
cyberBlocked := service.GetOpsCyberPolicy(c) != nil
|
||||
h.submitOpenAIUsageRecordTask(c.Request.Context(), result, func(ctx context.Context) {
|
||||
@@ -312,6 +313,7 @@ func (h *OpenAIGatewayHandler) ChatCompletions(c *gin.Context) {
|
||||
UserAgent: userAgent,
|
||||
IPAddress: clientIP,
|
||||
APIKeyService: h.apiKeyService,
|
||||
QuotaPlatform: quotaPlatform,
|
||||
ChannelUsageFields: channelMapping.ToUsageFields(reqModel, result.UpstreamModel),
|
||||
CyberBlocked: cyberBlocked,
|
||||
}); err != nil {
|
||||
|
||||
@@ -219,6 +219,7 @@ func (h *OpenAIGatewayHandler) Embeddings(c *gin.Context) {
|
||||
clientIP := ip.GetClientIP(c)
|
||||
inboundEndpoint := GetInboundEndpoint(c)
|
||||
upstreamEndpoint := GetUpstreamEndpoint(c, account.Platform)
|
||||
quotaPlatform := service.QuotaPlatform(c.Request.Context(), apiKey)
|
||||
|
||||
h.submitOpenAIUsageRecordTask(c.Request.Context(), result, func(ctx context.Context) {
|
||||
if err := h.gatewayService.RecordUsage(ctx, &service.OpenAIRecordUsageInput{
|
||||
@@ -232,6 +233,7 @@ func (h *OpenAIGatewayHandler) Embeddings(c *gin.Context) {
|
||||
UserAgent: userAgent,
|
||||
IPAddress: clientIP,
|
||||
APIKeyService: h.apiKeyService,
|
||||
QuotaPlatform: quotaPlatform,
|
||||
ChannelUsageFields: channelMapping.ToUsageFields(reqModel, result.UpstreamModel),
|
||||
}); err != nil {
|
||||
logger.L().With(
|
||||
|
||||
@@ -513,6 +513,7 @@ func (h *OpenAIGatewayHandler) Responses(c *gin.Context) {
|
||||
requestPayloadHash := service.HashUsageRequestPayload(body)
|
||||
inboundEndpoint := GetInboundEndpoint(c)
|
||||
upstreamEndpoint := resolveOpenAIUpstreamEndpoint(c, account)
|
||||
quotaPlatform := service.QuotaPlatform(c.Request.Context(), apiKey)
|
||||
|
||||
// 使用量记录通过有界 worker 池提交,避免请求热路径创建无界 goroutine。
|
||||
cyberBlocked := service.GetOpsCyberPolicy(c) != nil
|
||||
@@ -529,6 +530,7 @@ func (h *OpenAIGatewayHandler) Responses(c *gin.Context) {
|
||||
IPAddress: clientIP,
|
||||
RequestPayloadHash: requestPayloadHash,
|
||||
APIKeyService: h.apiKeyService,
|
||||
QuotaPlatform: quotaPlatform,
|
||||
ChannelUsageFields: channelMapping.ToUsageFields(reqModel, result.UpstreamModel),
|
||||
CyberBlocked: cyberBlocked,
|
||||
}); err != nil {
|
||||
@@ -926,6 +928,7 @@ func (h *OpenAIGatewayHandler) Messages(c *gin.Context) {
|
||||
requestPayloadHash := service.HashUsageRequestPayload(body)
|
||||
inboundEndpoint := GetInboundEndpoint(c)
|
||||
upstreamEndpoint := resolveOpenAIUpstreamEndpoint(c, account)
|
||||
quotaPlatform := service.QuotaPlatform(c.Request.Context(), apiKey)
|
||||
|
||||
cyberBlocked := service.GetOpsCyberPolicy(c) != nil
|
||||
h.submitOpenAIUsageRecordTask(c.Request.Context(), result, func(ctx context.Context) {
|
||||
@@ -941,6 +944,7 @@ func (h *OpenAIGatewayHandler) Messages(c *gin.Context) {
|
||||
IPAddress: clientIP,
|
||||
RequestPayloadHash: requestPayloadHash,
|
||||
APIKeyService: h.apiKeyService,
|
||||
QuotaPlatform: quotaPlatform,
|
||||
ChannelUsageFields: channelMappingMsg.ToUsageFields(reqModel, result.UpstreamModel),
|
||||
CyberBlocked: cyberBlocked,
|
||||
}); err != nil {
|
||||
@@ -1523,6 +1527,7 @@ func (h *OpenAIGatewayHandler) ResponsesWebSocket(c *gin.Context) {
|
||||
h.gatewayService.ReportOpenAIAccountScheduleResult(account.ID, true, result.FirstTokenMs)
|
||||
inboundEndpoint := GetInboundEndpoint(c)
|
||||
upstreamEndpoint := resolveOpenAIUpstreamEndpoint(c, account)
|
||||
quotaPlatform := service.QuotaPlatform(c.Request.Context(), apiKey)
|
||||
cyberBlocked := service.GetOpsCyberPolicy(c) != nil
|
||||
h.submitOpenAIUsageRecordTask(ctx, result, func(taskCtx context.Context) {
|
||||
if err := h.gatewayService.RecordUsage(taskCtx, &service.OpenAIRecordUsageInput{
|
||||
@@ -1537,6 +1542,7 @@ func (h *OpenAIGatewayHandler) ResponsesWebSocket(c *gin.Context) {
|
||||
IPAddress: clientIP,
|
||||
RequestPayloadHash: requestPayloadHash,
|
||||
APIKeyService: h.apiKeyService,
|
||||
QuotaPlatform: quotaPlatform,
|
||||
ChannelUsageFields: channelMappingWS.ToUsageFields(reqModel, result.UpstreamModel),
|
||||
CyberBlocked: cyberBlocked,
|
||||
}); err != nil {
|
||||
|
||||
@@ -341,6 +341,7 @@ func (h *OpenAIGatewayHandler) Images(c *gin.Context) {
|
||||
}
|
||||
inboundEndpoint := GetInboundEndpoint(c)
|
||||
upstreamEndpoint := GetUpstreamEndpoint(c, account.Platform)
|
||||
quotaPlatform := service.QuotaPlatform(c.Request.Context(), apiKey)
|
||||
|
||||
upstreamModel := ""
|
||||
if result != nil {
|
||||
@@ -359,6 +360,7 @@ func (h *OpenAIGatewayHandler) Images(c *gin.Context) {
|
||||
IPAddress: clientIP,
|
||||
RequestPayloadHash: requestPayloadHash,
|
||||
APIKeyService: h.apiKeyService,
|
||||
QuotaPlatform: quotaPlatform,
|
||||
ChannelUsageFields: channelMapping.ToUsageFields(requestModel, upstreamModel),
|
||||
}); err != nil {
|
||||
logger.L().With(
|
||||
|
||||
@@ -0,0 +1,65 @@
|
||||
package handler
|
||||
|
||||
import (
|
||||
"go/ast"
|
||||
"go/parser"
|
||||
"go/token"
|
||||
"path/filepath"
|
||||
"testing"
|
||||
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
func TestOpenAIRecordUsageInputsCarryQuotaPlatform(t *testing.T) {
|
||||
files := []string{
|
||||
"openai_gateway_handler.go",
|
||||
"openai_chat_completions.go",
|
||||
"openai_embeddings.go",
|
||||
"openai_images.go",
|
||||
}
|
||||
|
||||
for _, name := range files {
|
||||
t.Run(name, func(t *testing.T) {
|
||||
fset := token.NewFileSet()
|
||||
file, err := parser.ParseFile(fset, filepath.Join(".", name), nil, 0)
|
||||
require.NoError(t, err)
|
||||
|
||||
var missing []token.Position
|
||||
ast.Inspect(file, func(node ast.Node) bool {
|
||||
literal, ok := node.(*ast.CompositeLit)
|
||||
if !ok || !isOpenAIRecordUsageInputLiteral(literal.Type) {
|
||||
return true
|
||||
}
|
||||
if !compositeLiteralHasKey(literal, "QuotaPlatform") {
|
||||
missing = append(missing, fset.Position(literal.Lbrace))
|
||||
}
|
||||
return true
|
||||
})
|
||||
|
||||
require.Empty(t, missing, "OpenAI usage post-billing must receive request-time QuotaPlatform")
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func isOpenAIRecordUsageInputLiteral(expr ast.Expr) bool {
|
||||
selector, ok := expr.(*ast.SelectorExpr)
|
||||
if !ok {
|
||||
return false
|
||||
}
|
||||
pkg, ok := selector.X.(*ast.Ident)
|
||||
return ok && pkg.Name == "service" && selector.Sel.Name == "OpenAIRecordUsageInput"
|
||||
}
|
||||
|
||||
func compositeLiteralHasKey(literal *ast.CompositeLit, key string) bool {
|
||||
for _, elt := range literal.Elts {
|
||||
pair, ok := elt.(*ast.KeyValueExpr)
|
||||
if !ok {
|
||||
continue
|
||||
}
|
||||
ident, ok := pair.Key.(*ast.Ident)
|
||||
if ok && ident.Name == key {
|
||||
return true
|
||||
}
|
||||
}
|
||||
return false
|
||||
}
|
||||
@@ -6096,6 +6096,7 @@ type OpenAIRecordUsageInput struct {
|
||||
IPAddress string // 请求的客户端 IP 地址
|
||||
RequestPayloadHash string
|
||||
APIKeyService APIKeyQuotaUpdater
|
||||
QuotaPlatform string // user×platform quota platform resolved by the handler before async billing.
|
||||
// CyberBlocked 为 true 时把该用量行标记为 cyber(request_type=cyber),计费逻辑不变。
|
||||
CyberBlocked bool
|
||||
ChannelUsageFields
|
||||
@@ -6370,6 +6371,13 @@ func (s *OpenAIGatewayService) RecordUsage(ctx context.Context, input *OpenAIRec
|
||||
return nil
|
||||
}
|
||||
|
||||
// Async usage billing runs outside the original request context, so it
|
||||
// cannot recover ForcePlatform there. Fall back for internal/test callers.
|
||||
quotaPlatform := input.QuotaPlatform
|
||||
if quotaPlatform == "" {
|
||||
quotaPlatform = PlatformFromAPIKey(apiKey)
|
||||
}
|
||||
|
||||
billingErr := func() error {
|
||||
_, err := applyUsageBilling(ctx, requestID, usageLog, &postUsageBillingParams{
|
||||
Cost: cost,
|
||||
@@ -6381,7 +6389,7 @@ func (s *OpenAIGatewayService) RecordUsage(ctx context.Context, input *OpenAIRec
|
||||
IsSubscriptionBill: isSubscriptionBilling,
|
||||
AccountRateMultiplier: accountRateMultiplier,
|
||||
APIKeyService: input.APIKeyService,
|
||||
Platform: PlatformFromAPIKey(apiKey),
|
||||
Platform: quotaPlatform,
|
||||
}, s.billingDeps(), s.usageBillingRepo)
|
||||
return err
|
||||
}()
|
||||
|
||||
Reference in New Issue
Block a user