feat(antigravity): 支持按映射模型计费

This commit is contained in:
song
2026-01-20 10:59:10 +08:00
parent d1a6303e49
commit c43aa22cdb

View File

@@ -29,7 +29,10 @@ const (
antigravityRetryMaxDelay = 16 * time.Second
)
const antigravityScopeRateLimitEnv = "GATEWAY_ANTIGRAVITY_429_SCOPE_LIMIT"
const (
antigravityScopeRateLimitEnv = "GATEWAY_ANTIGRAVITY_429_SCOPE_LIMIT"
antigravityBillingModelEnv = "GATEWAY_ANTIGRAVITY_BILL_WITH_MAPPED_MODEL"
)
// antigravityRetryLoopParams 重试循环的参数
type antigravityRetryLoopParams struct {
@@ -711,6 +714,10 @@ func (s *AntigravityGatewayService) Forward(ctx context.Context, c *gin.Context,
originalModel := claudeReq.Model
mappedModel := s.getMappedModel(account, claudeReq.Model)
quotaScope, _ := resolveAntigravityQuotaScope(originalModel)
billingModel := originalModel
if antigravityUseMappedModelForBilling() && strings.TrimSpace(mappedModel) != "" {
billingModel = mappedModel
}
// 获取 access_token
if s.tokenProvider == nil {
@@ -971,7 +978,7 @@ func (s *AntigravityGatewayService) Forward(ctx context.Context, c *gin.Context,
return &ForwardResult{
RequestID: requestID,
Usage: *usage,
Model: originalModel, // 使用原始模型用于计费和日志
Model: billingModel,
Stream: claudeReq.Stream,
Duration: time.Since(startTime),
FirstTokenMs: firstTokenMs,
@@ -1280,6 +1287,10 @@ func (s *AntigravityGatewayService) ForwardGemini(ctx context.Context, c *gin.Co
}
mappedModel := s.getMappedModel(account, originalModel)
billingModel := originalModel
if antigravityUseMappedModelForBilling() && strings.TrimSpace(mappedModel) != "" {
billingModel = mappedModel
}
// 获取 access_token
if s.tokenProvider == nil {
@@ -1478,7 +1489,7 @@ handleSuccess:
return &ForwardResult{
RequestID: requestID,
Usage: *usage,
Model: originalModel,
Model: billingModel,
Stream: stream,
Duration: time.Since(startTime),
FirstTokenMs: firstTokenMs,
@@ -1525,6 +1536,11 @@ func antigravityUseScopeRateLimit() bool {
return v == "1" || v == "true" || v == "yes" || v == "on"
}
func antigravityUseMappedModelForBilling() bool {
v := strings.ToLower(strings.TrimSpace(os.Getenv(antigravityBillingModelEnv)))
return v == "1" || v == "true" || v == "yes" || v == "on"
}
func (s *AntigravityGatewayService) handleUpstreamError(ctx context.Context, prefix string, account *Account, statusCode int, headers http.Header, body []byte, quotaScope AntigravityQuotaScope) {
// 429 使用 Gemini 格式解析(从 body 解析重置时间)
if statusCode == 429 {