coai/auth/plan.go

package auth

import (
	"chat/globals"
	"chat/utils"
	"database/sql"
	"github.com/go-redis/redis/v8"
)

type Plan struct {
	Level int
	Price float32
	Usage []PlanUsage
}

type PlanUsage struct {
	Id        string
	Value     int64
	Including func(string) bool
}

type Usage struct {
	Used  int64 `json:"used"`
	Total int64 `json:"total"`
}
type UsageMap map[string]Usage

var Plans = []Plan{
	{
		Level: 0,
		Price: 0,
		Usage: []PlanUsage{},
	},
	{
		Level: 1,
		Price: 18,
		Usage: []PlanUsage{
			{Id: "gpt-4", Value: 25, Including: globals.IsGPT4NativeModel},
			{Id: "claude-100k", Value: 50, Including: globals.IsClaude100KModel},
		},
	},
	{
		Level: 2,
		Price: 36,
		Usage: []PlanUsage{
			{Id: "gpt-4", Value: 50, Including: globals.IsGPT4NativeModel},
			{Id: "claude-100k", Value: 100, Including: globals.IsClaude100KModel},
			{Id: "midjourney", Value: 25, Including: globals.IsMidjourneyFastModel},
		},
	},
	{
		Level: 3,
		Price: 72,
		Usage: []PlanUsage{
			{Id: "gpt-4", Value: 100, Including: globals.IsGPT4NativeModel},
			{Id: "claude-100k", Value: 200, Including: globals.IsClaude100KModel},
			{Id: "midjourney", Value: 50, Including: globals.IsMidjourneyFastModel},
		},
	},
	{
		// enterprise
		Level: 4,
		Price: 999,
		Usage: []PlanUsage{},
	},
}

func IncreaseSubscriptionUsage(cache *redis.Client, user *User, t string, limit int64) bool {
	return utils.IncrWithLimit(cache, globals.GetSubscriptionLimitFormat(t, user.ID), 1, limit, 60*60*24) // 1 day
}

func DecreaseSubscriptionUsage(cache *redis.Client, user *User, t string) bool {
	return utils.DecrInt(cache, globals.GetSubscriptionLimitFormat(t, user.ID), 1)
}

func (p *Plan) GetUsage(user *User, db *sql.DB, cache *redis.Client) UsageMap {
	return utils.EachObject[PlanUsage, Usage](p.Usage, func(usage PlanUsage) (string, Usage) {
		return usage.Id, usage.GetUsageForm(user, db, cache)
	})
}

func (p *PlanUsage) GetUsage(user *User, db *sql.DB, cache *redis.Client) int64 {
	return utils.MustInt(cache, globals.GetSubscriptionLimitFormat(p.Id, user.GetID(db)))
}

func (p *PlanUsage) GetUsageForm(user *User, db *sql.DB, cache *redis.Client) Usage {
	return Usage{
		Used:  p.GetUsage(user, db, cache),
		Total: p.Value,
	}
}

func (p *PlanUsage) IsInfinity() bool {
	return p.Value == -1
}

func (p *PlanUsage) IsExceeded(user *User, db *sql.DB, cache *redis.Client) bool {
	return p.IsInfinity() || p.GetUsage(user, db, cache) < p.Value
}

func (p *PlanUsage) Increase(user *User, cache *redis.Client) bool {
	if p.Value == -1 {
		return true
	}
	return IncreaseSubscriptionUsage(cache, user, p.Id, p.Value)
}

func (p *PlanUsage) Decrease(user *User, cache *redis.Client) bool {
	if p.Value == -1 {
		return true
	}
	return DecreaseSubscriptionUsage(cache, user, p.Id)
}

func (u *User) GetSubscriptionUsage(db *sql.DB, cache *redis.Client) UsageMap {
	plan := u.GetPlan(db)
	return plan.GetUsage(u, db, cache)
}

func (p *Plan) IncreaseUsage(user *User, cache *redis.Client, model string) bool {
	for _, usage := range p.Usage {
		if usage.Including(model) {
			return usage.Increase(user, cache)
		}
	}

	return false
}

func (p *Plan) DecreaseUsage(user *User, cache *redis.Client, model string) bool {
	for _, usage := range p.Usage {
		if usage.Including(model) {
			return usage.Decrease(user, cache)
		}
	}

	return false
}

func GetLevel(level int) Plan {
	for _, plan := range Plans {
		if plan.Level == level {
			return plan
		}
	}
	return Plan{}
}

func InLevel(level int) bool {
	return utils.InRange(level, 1, 3)
}