Browse Source

feat: add AI billing module (quota + usage)

- QuotaService: CheckAndFreeze, Settle, Unfreeze
- UsageService: Record, UpdateConversationStats
- Skip billing for user-provided API keys
master
dark 1 month ago
parent
commit
0e3173189f
  1. 50
      backend/internal/ai/billing/quota.go
  2. 32
      backend/internal/ai/billing/usage.go

50
backend/internal/ai/billing/quota.go

@ -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)
}

32
backend/internal/ai/billing/usage.go

@ -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…
Cancel
Save