Skip to content

Commit

Permalink
feat: /v1/models now only return available models
Browse files Browse the repository at this point in the history
  • Loading branch information
songquanpeng committed Apr 3, 2024
1 parent dc7aaf2 commit 8b9813d
Show file tree
Hide file tree
Showing 2 changed files with 36 additions and 2 deletions.
35 changes: 34 additions & 1 deletion controller/model.go
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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,
})
}

Expand Down
3 changes: 2 additions & 1 deletion middleware/auth.go
Original file line number Diff line number Diff line change
Expand Up @@ -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
}
Expand Down

0 comments on commit 8b9813d

Please sign in to comment.