Browse Source
- QuotaService: CheckAndFreeze, Settle, Unfreeze - UsageService: Record, UpdateConversationStats - Skip billing for user-provided API keysmaster
2 changed files with 82 additions and 0 deletions
@ -0,0 +1,50 @@ |
|||||
|
package billing |
||||
|
|
||||
|
import ( |
||||
|
"context" |
||||
|
|
||||
|
"github.com/youruser/base/model" |
||||
|
|
||||
|
"gorm.io/gorm" |
||||
|
) |
||||
|
|
||||
|
// QuotaService handles quota freeze/settle/unfreeze operations
|
||||
|
type QuotaService struct{} |
||||
|
|
||||
|
func NewQuotaService() *QuotaService { |
||||
|
return &QuotaService{} |
||||
|
} |
||||
|
|
||||
|
// CheckAndFreeze checks if user has sufficient balance and freezes the estimated cost.
|
||||
|
// Returns nil if successful, error if insufficient balance.
|
||||
|
// If apiKeyId > 0, it's a user-provided key — skip billing.
|
||||
|
func (s *QuotaService) CheckAndFreeze(ctx context.Context, db *gorm.DB, userId int64, estimatedCost float64, apiKeyId int64) error { |
||||
|
// User-provided keys skip billing
|
||||
|
if apiKeyId > 0 { |
||||
|
return nil |
||||
|
} |
||||
|
|
||||
|
// Ensure quota record exists
|
||||
|
_, err := model.AIUserQuotaEnsure(ctx, db, userId) |
||||
|
if err != nil { |
||||
|
return err |
||||
|
} |
||||
|
|
||||
|
return model.AIUserQuotaFreeze(ctx, db, userId, estimatedCost) |
||||
|
} |
||||
|
|
||||
|
// Settle finalizes billing: releases frozen amount, deducts actual cost, refunds difference.
|
||||
|
func (s *QuotaService) Settle(ctx context.Context, db *gorm.DB, userId int64, frozenAmount, actualCost float64, apiKeyId int64) error { |
||||
|
if apiKeyId > 0 { |
||||
|
return nil |
||||
|
} |
||||
|
return model.AIUserQuotaSettle(ctx, db, userId, frozenAmount, actualCost) |
||||
|
} |
||||
|
|
||||
|
// Unfreeze releases frozen amount back to balance (used on error).
|
||||
|
func (s *QuotaService) Unfreeze(ctx context.Context, db *gorm.DB, userId int64, amount float64, apiKeyId int64) error { |
||||
|
if apiKeyId > 0 { |
||||
|
return nil |
||||
|
} |
||||
|
return model.AIUserQuotaUnfreeze(ctx, db, userId, amount) |
||||
|
} |
||||
@ -0,0 +1,32 @@ |
|||||
|
package billing |
||||
|
|
||||
|
import ( |
||||
|
"context" |
||||
|
|
||||
|
"github.com/youruser/base/model" |
||||
|
|
||||
|
"gorm.io/gorm" |
||||
|
) |
||||
|
|
||||
|
// UsageService handles usage recording
|
||||
|
type UsageService struct{} |
||||
|
|
||||
|
func NewUsageService() *UsageService { |
||||
|
return &UsageService{} |
||||
|
} |
||||
|
|
||||
|
// Record inserts a usage record
|
||||
|
func (s *UsageService) Record(ctx context.Context, db *gorm.DB, record *model.AIUsageRecord) error { |
||||
|
_, err := model.AIUsageRecordInsert(ctx, db, record) |
||||
|
return err |
||||
|
} |
||||
|
|
||||
|
// UpdateConversationStats updates conversation token count and cost
|
||||
|
func (s *UsageService) UpdateConversationStats(ctx context.Context, db *gorm.DB, conversationId int64, tokens int64, cost float64) error { |
||||
|
return db.WithContext(ctx).Model(&model.AIConversation{}). |
||||
|
Where("id = ?", conversationId). |
||||
|
Updates(map[string]interface{}{ |
||||
|
"total_tokens": gorm.Expr("total_tokens + ?", tokens), |
||||
|
"total_cost": gorm.Expr("total_cost + ?", cost), |
||||
|
}).Error |
||||
|
} |
||||
Loading…
Reference in new issue