feat: /v1/models now only return available models
This commit is contained in:
parent
dc7aaf2de5
commit
8b9813d63b
@ -11,6 +11,7 @@ import (
|
||||
relaymodel "github.com/songquanpeng/one-api/relay/model"
|
||||
"github.com/songquanpeng/one-api/relay/util"
|
||||
"net/http"
|
||||
"strings"
|
||||
)
|
||||
|
||||
// https://platform.openai.com/docs/api-reference/models/list
|
||||
@ -121,9 +122,41 @@ func DashboardListModels(c *gin.Context) {
|
||||
}
|
||||
|
||||
func ListModels(c *gin.Context) {
|
||||
ctx := c.Request.Context()
|
||||
var availableModels []string
|
||||
if c.GetString("available_models") != "" {
|
||||
availableModels = strings.Split(c.GetString("available_models"), ",")
|
||||
} else {
|
||||
userId := c.GetInt("id")
|
||||
userGroup, _ := model.CacheGetUserGroup(userId)
|
||||
availableModels, _ = model.CacheGetGroupModels(ctx, userGroup)
|
||||
}
|
||||
modelSet := make(map[string]bool)
|
||||
for _, availableModel := range availableModels {
|
||||
modelSet[availableModel] = true
|
||||
}
|
||||
var availableOpenAIModels []OpenAIModels
|
||||
for _, model := range openAIModels {
|
||||
if _, ok := modelSet[model.Id]; ok {
|
||||
modelSet[model.Id] = false
|
||||
availableOpenAIModels = append(availableOpenAIModels, model)
|
||||
}
|
||||
}
|
||||
for modelName, ok := range modelSet {
|
||||
if ok {
|
||||
availableOpenAIModels = append(availableOpenAIModels, OpenAIModels{
|
||||
Id: modelName,
|
||||
Object: "model",
|
||||
Created: 1626777600,
|
||||
OwnedBy: "custom",
|
||||
Root: modelName,
|
||||
Parent: nil,
|
||||
})
|
||||
}
|
||||
}
|
||||
c.JSON(200, gin.H{
|
||||
"object": "list",
|
||||
"data": openAIModels,
|
||||
"data": availableOpenAIModels,
|
||||
})
|
||||
}
|
||||
|
||||
|
@ -115,7 +115,8 @@ func TokenAuth() func(c *gin.Context) {
|
||||
}
|
||||
c.Set("request_model", requestModel)
|
||||
if token.Models != nil && *token.Models != "" {
|
||||
if !isModelInList(requestModel, *token.Models) {
|
||||
c.Set("available_models", *token.Models)
|
||||
if requestModel != "" && !isModelInList(requestModel, *token.Models) {
|
||||
abortWithMessage(c, http.StatusForbidden, fmt.Sprintf("该令牌无权使用模型:%s", requestModel))
|
||||
return
|
||||
}
|
||||
|
Loading…
Reference in New Issue
Block a user