66 lines
1.9 KiB
Go
66 lines
1.9 KiB
Go
package controller
|
|
|
|
import (
|
|
"github.com/QuantumNous/new-api/common"
|
|
"github.com/QuantumNous/new-api/constant"
|
|
"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" || policy.Mode == "smart" {
|
|
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
|
|
}
|
|
common.SetContextKey(c, constant.ContextKeyUserGroup, group)
|
|
c.Set(model.TokenRoutingContextKey, &request.Policy)
|
|
entries := service.ResolveRequestPool(c, request.Model)
|
|
common.ApiSuccess(c, entries)
|
|
}
|