feat: able to set model limitation for token (close #178)

model/channel.go

@@ -1,6 +1,7 @@
 package model
 
 import (
+	"context"
 	"encoding/json"
 	"fmt"
 	"github.com/songquanpeng/one-api/common"
@@ -8,6 +9,8 @@ import (
 	"github.com/songquanpeng/one-api/common/helper"
 	"github.com/songquanpeng/one-api/common/logger"
 	"gorm.io/gorm"
+	"sort"
+	"strings"
 )
 
 type Channel struct {
@@ -25,7 +28,7 @@ type Channel struct {
 	Balance            float64 `json:"balance"` // in USD
 	BalanceUpdatedTime int64   `json:"balance_updated_time" gorm:"bigint"`
 	Models             string  `json:"models"`
-	Group              string  `json:"group" gorm:"type:varchar(32);default:'default'"`
+	Group              string  `json:"group" gorm:"index;type:varchar(32);default:'default'"`
 	UsedQuota          int64   `json:"used_quota" gorm:"bigint;default:0"`
 	ModelMapping       *string `json:"model_mapping" gorm:"type:varchar(1024);default:''"`
 	Priority           *int64  `json:"priority" gorm:"bigint;default:0"`
@@ -202,3 +205,28 @@ func DeleteDisabledChannel() (int64, error) {
 	result := DB.Where("status = ? or status = ?", common.ChannelStatusAutoDisabled, common.ChannelStatusManuallyDisabled).Delete(&Channel{})
 	return result.RowsAffected, result.Error
 }
+
+func GetGroupModels(ctx context.Context, group string) ([]string, error) {
+	groupCol := "`group`"
+	if common.UsingPostgreSQL {
+		groupCol = `"group"`
+	}
+	var modelsList []string
+	err := DB.Model(&Channel{}).Distinct("models").Where(groupCol+" = ?", group).Pluck("models", &modelsList).Error
+	if err != nil {
+		return nil, err
+	}
+	set := make(map[string]bool)
+	for i := 0; i < len(modelsList); i++ {
+		modelList := strings.Split(modelsList[i], ",")
+		for _, model := range modelList {
+			set[model] = true
+		}
+	}
+	modelList := make([]string, 0, len(set))
+	for model := range set {
+		modelList = append(modelList, model)
+	}
+	sort.Strings(modelList)
+	return modelList, err
+}