new-api-custom/controller/token_routing.go

package controller

import (
	"github.com/QuantumNous/new-api/common"
	"github.com/QuantumNous/new-api/model"
	"github.com/QuantumNous/new-api/service"
	"github.com/QuantumNous/new-api/setting"
	"github.com/gin-gonic/gin"
)

func setTokenRoutingPolicy(c *gin.Context, token *model.Token, policy *model.TokenRoutingPolicy) bool {
	body, err := common.Marshal(policy)
	if err != nil {
		common.ApiError(c, err)
		return false
	}
	if token.RoutingPolicy != "" && token.RoutingPolicy == string(body) {
		return true
	}
	userGroup, err := getTokenRequestUserGroup(c)
	if err != nil {
		common.ApiError(c, err)
		return false
	}
	if err := policy.Validate(setting.GetMaxTokenAutoGroups()+1, func(group string) bool { return service.IsUserSelectableGroup(userGroup, group) }); err != nil {
		common.ApiError(c, err)
		return false
	}
	token.RoutingPolicy = string(body)
	token.CrossGroupRetry = true
	token.AutoGroups = ""
	if policy.Mode == "inherit" {
		token.Group = "auto"
	} else if policy.Mode == "user" {
		token.Group = ""
	} else {
		token.Group = policy.Resolve("", nil, func(string) bool { return true })[0].Group
	}
	return true
}

func PreviewTokenRouting(c *gin.Context) {
	var request struct {
		Policy model.TokenRoutingPolicy `json:"routing_policy"`
		Model  string                   `json:"model"`
	}
	if err := c.ShouldBindJSON(&request); err != nil {
		common.ApiError(c, err)
		return
	}
	if err := request.Policy.Validate(setting.GetMaxTokenAutoGroups()+1, func(string) bool { return true }); err != nil {
		common.ApiError(c, err)
		return
	}
	group, err := getTokenRequestUserGroup(c)
	if err != nil {
		common.ApiError(c, err)
		return
	}
	inherited := service.GetUserAutoGroup(group)
	if request.Policy.Mode == "user" {
		inherited = []string{group}
	}
	entries := request.Policy.Resolve(request.Model, inherited, func(candidate string) bool { return service.IsUserSelectableGroup(group, candidate) })
	common.ApiSuccess(c, entries)
}