coai/auth/rule.go

package auth

import (
	"chat/channel"
	"database/sql"
	"fmt"

	"chat/globals"
	"chat/utils"

	"github.com/go-redis/redis/v8"
)

const (
	ErrNotAuthenticated = "not authenticated error (model: %s)"
	ErrNotSetPrice      = "the price of the model is not set (model: %s)"
	ErrNotEnoughQuota   = "user quota is not enough error (model: %s, minimum quota: %0.2f, your quota: %0.2f)"
	ErrEstimatedCost    = "estimated cost exceeds user quota (model: %s, estimated cost: %0.2f, your quota: %0.2f)"
)

// CanEnableModel returns whether the model can be enabled (without subscription)
func CanEnableModel(db *sql.DB, user *User, model string, messages []globals.Message) error {
	isAuth := user != nil
	isAdmin := isAuth && user.IsAdmin(db)

	charge := channel.ChargeInstance.GetCharge(model)

	if charge.IsUnsetType() && !isAdmin {
		return fmt.Errorf(ErrNotSetPrice, model)
	}

	if !charge.IsBilling() {
		// return if is the user is authenticated or anonymous is allowed for this model
		if charge.SupportAnonymous() || isAuth {
			return nil
		}

		return fmt.Errorf(ErrNotAuthenticated, model)
	}

	if !isAuth {
		return fmt.Errorf(ErrNotAuthenticated, model)
	}

	// Calculate estimated input cost
	inputTokens := utils.NumTokensFromMessages(messages, model, false)
	estimatedInputCost := float32(inputTokens) / 1000 * charge.GetInput()

	// Get user's current quota
	quota := user.GetQuota(db)
	if quota < estimatedInputCost {
		return fmt.Errorf(ErrEstimatedCost, model, estimatedInputCost, quota)
	}

	return nil
}

func CanEnableModelWithSubscription(db *sql.DB, cache *redis.Client, user *User, model string, messages []globals.Message) (canEnable error, usePlan bool) {
	// use subscription quota first
	if user != nil && HandleSubscriptionUsage(db, cache, user, model) {
		return nil, true
	}
	return CanEnableModel(db, user, model, messages), false
}