From 72983ac734859cb18fcd50fd0f8a0372b5fdcb6a Mon Sep 17 00:00:00 2001 From: Martial BE Date: Thu, 28 Dec 2023 15:48:20 +0800 Subject: [PATCH] =?UTF-8?q?=E2=9C=A8=20Model=20list=20return=20by=20group?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- common/image/image.go | 2 +- controller/model.go | 602 ++++------------------------------------- model/ability.go | 16 ++ model/cache.go | 19 ++ router/api-router.go | 2 +- router/relay-router.go | 2 +- 6 files changed, 95 insertions(+), 548 deletions(-) diff --git a/common/image/image.go b/common/image/image.go index f4998bc9..aff8cb31 100644 --- a/common/image/image.go +++ b/common/image/image.go @@ -63,7 +63,7 @@ func GetImageFromUrl(url string) (mimeType string, data string, err error) { isImage, err := IsImageUrl(url) if !isImage { if err == nil { - err = errors.New("Invalid image link") + err = errors.New("invalid image link") } return } diff --git a/controller/model.go b/controller/model.go index 31dc06d1..9e023c93 100644 --- a/controller/model.go +++ b/controller/model.go @@ -2,7 +2,11 @@ package controller import ( "fmt" + "net/http" + "one-api/common" + "one-api/model" "one-api/types" + "sort" "github.com/gin-gonic/gin" ) @@ -25,559 +29,38 @@ type OpenAIModelPermission struct { } type OpenAIModels struct { - Id string `json:"id"` - Object string `json:"object"` - Created int `json:"created"` - OwnedBy string `json:"owned_by"` - Permission []OpenAIModelPermission `json:"permission"` - Root string `json:"root"` - Parent *string `json:"parent"` + Id string `json:"id"` + Object string `json:"object"` + Created int `json:"created"` + OwnedBy *string `json:"owned_by"` + Permission *[]OpenAIModelPermission `json:"permission"` + Root *string `json:"root"` + Parent *string `json:"parent"` } var openAIModels []OpenAIModels var openAIModelsMap map[string]OpenAIModels func init() { - var permission []OpenAIModelPermission - permission = append(permission, OpenAIModelPermission{ - Id: "modelperm-LwHkVFn8AcMItP432fKKDIKJ", - Object: "model_permission", - Created: 1626777600, - AllowCreateEngine: true, - AllowSampling: true, - AllowLogprobs: true, - AllowSearchIndices: false, - AllowView: true, - AllowFineTuning: false, - Organization: "*", - Group: nil, - IsBlocking: false, - }) // https://platform.openai.com/docs/models/model-endpoint-compatibility - openAIModels = []OpenAIModels{ - { - Id: "dall-e-2", - Object: "model", - Created: 1677649963, - OwnedBy: "openai", - Permission: permission, - Root: "dall-e-2", - Parent: nil, - }, - { - Id: "dall-e-3", - Object: "model", - Created: 1677649963, - OwnedBy: "openai", - Permission: permission, - Root: "dall-e-3", - Parent: nil, - }, - { - Id: "whisper-1", - Object: "model", - Created: 1677649963, - OwnedBy: "openai", - Permission: permission, - Root: "whisper-1", - Parent: nil, - }, - { - Id: "tts-1", - Object: "model", - Created: 1677649963, - OwnedBy: "openai", - Permission: permission, - Root: "tts-1", - Parent: nil, - }, - { - Id: "tts-1-1106", - Object: "model", - Created: 1677649963, - OwnedBy: "openai", - Permission: permission, - Root: "tts-1-1106", - Parent: nil, - }, - { - Id: "tts-1-hd", - Object: "model", - Created: 1677649963, - OwnedBy: "openai", - Permission: permission, - Root: "tts-1-hd", - Parent: nil, - }, - { - Id: "tts-1-hd-1106", - Object: "model", - Created: 1677649963, - OwnedBy: "openai", - Permission: permission, - Root: "tts-1-hd-1106", - Parent: nil, - }, - { - Id: "gpt-3.5-turbo", - Object: "model", - Created: 1677649963, - OwnedBy: "openai", - Permission: permission, - Root: "gpt-3.5-turbo", - Parent: nil, - }, - { - Id: "gpt-3.5-turbo-0301", - Object: "model", - Created: 1677649963, - OwnedBy: "openai", - Permission: permission, - Root: "gpt-3.5-turbo-0301", - Parent: nil, - }, - { - Id: "gpt-3.5-turbo-0613", - Object: "model", - Created: 1677649963, - OwnedBy: "openai", - Permission: permission, - Root: "gpt-3.5-turbo-0613", - Parent: nil, - }, - { - Id: "gpt-3.5-turbo-16k", - Object: "model", - Created: 1677649963, - OwnedBy: "openai", - Permission: permission, - Root: "gpt-3.5-turbo-16k", - Parent: nil, - }, - { - Id: "gpt-3.5-turbo-16k-0613", - Object: "model", - Created: 1677649963, - OwnedBy: "openai", - Permission: permission, - Root: "gpt-3.5-turbo-16k-0613", - Parent: nil, - }, - { - Id: "gpt-3.5-turbo-1106", - Object: "model", - Created: 1699593571, - OwnedBy: "openai", - Permission: permission, - Root: "gpt-3.5-turbo-1106", - Parent: nil, - }, - { - Id: "gpt-3.5-turbo-instruct", - Object: "model", - Created: 1677649963, - OwnedBy: "openai", - Permission: permission, - Root: "gpt-3.5-turbo-instruct", - Parent: nil, - }, - { - Id: "gpt-4", - Object: "model", - Created: 1677649963, - OwnedBy: "openai", - Permission: permission, - Root: "gpt-4", - Parent: nil, - }, - { - Id: "gpt-4-0314", - Object: "model", - Created: 1677649963, - OwnedBy: "openai", - Permission: permission, - Root: "gpt-4-0314", - Parent: nil, - }, - { - Id: "gpt-4-0613", - Object: "model", - Created: 1677649963, - OwnedBy: "openai", - Permission: permission, - Root: "gpt-4-0613", - Parent: nil, - }, - { - Id: "gpt-4-32k", - Object: "model", - Created: 1677649963, - OwnedBy: "openai", - Permission: permission, - Root: "gpt-4-32k", - Parent: nil, - }, - { - Id: "gpt-4-32k-0314", - Object: "model", - Created: 1677649963, - OwnedBy: "openai", - Permission: permission, - Root: "gpt-4-32k-0314", - Parent: nil, - }, - { - Id: "gpt-4-32k-0613", - Object: "model", - Created: 1677649963, - OwnedBy: "openai", - Permission: permission, - Root: "gpt-4-32k-0613", - Parent: nil, - }, - { - Id: "gpt-4-1106-preview", - Object: "model", - Created: 1699593571, - OwnedBy: "openai", - Permission: permission, - Root: "gpt-4-1106-preview", - Parent: nil, - }, - { - Id: "gpt-4-vision-preview", - Object: "model", - Created: 1699593571, - OwnedBy: "openai", - Permission: permission, - Root: "gpt-4-vision-preview", - Parent: nil, - }, - { - Id: "text-embedding-ada-002", - Object: "model", - Created: 1677649963, - OwnedBy: "openai", - Permission: permission, - Root: "text-embedding-ada-002", - Parent: nil, - }, - { - Id: "text-davinci-003", - Object: "model", - Created: 1677649963, - OwnedBy: "openai", - Permission: permission, - Root: "text-davinci-003", - Parent: nil, - }, - { - Id: "text-davinci-002", - Object: "model", - Created: 1677649963, - OwnedBy: "openai", - Permission: permission, - Root: "text-davinci-002", - Parent: nil, - }, - { - Id: "text-curie-001", - Object: "model", - Created: 1677649963, - OwnedBy: "openai", - Permission: permission, - Root: "text-curie-001", - Parent: nil, - }, - { - Id: "text-babbage-001", - Object: "model", - Created: 1677649963, - OwnedBy: "openai", - Permission: permission, - Root: "text-babbage-001", - Parent: nil, - }, - { - Id: "text-ada-001", - Object: "model", - Created: 1677649963, - OwnedBy: "openai", - Permission: permission, - Root: "text-ada-001", - Parent: nil, - }, - { - Id: "text-moderation-latest", - Object: "model", - Created: 1677649963, - OwnedBy: "openai", - Permission: permission, - Root: "text-moderation-latest", - Parent: nil, - }, - { - Id: "text-moderation-stable", - Object: "model", - Created: 1677649963, - OwnedBy: "openai", - Permission: permission, - Root: "text-moderation-stable", - Parent: nil, - }, - { - Id: "text-davinci-edit-001", - Object: "model", - Created: 1677649963, - OwnedBy: "openai", - Permission: permission, - Root: "text-davinci-edit-001", - Parent: nil, - }, - { - Id: "code-davinci-edit-001", - Object: "model", - Created: 1677649963, - OwnedBy: "openai", - Permission: permission, - Root: "code-davinci-edit-001", - Parent: nil, - }, - { - Id: "claude-instant-1", - Object: "model", - Created: 1677649963, - OwnedBy: "anthropic", - Permission: permission, - Root: "claude-instant-1", - Parent: nil, - }, - { - Id: "claude-2", - Object: "model", - Created: 1677649963, - OwnedBy: "anthropic", - Permission: permission, - Root: "claude-2", - Parent: nil, - }, - { - Id: "claude-2.1", - Object: "model", - Created: 1677649963, - OwnedBy: "anthropic", - Permission: permission, - Root: "claude-2.1", - Parent: nil, - }, - { - Id: "claude-2.0", - Object: "model", - Created: 1677649963, - OwnedBy: "anthropic", - Permission: permission, - Root: "claude-2.0", - Parent: nil, - }, - { - Id: "ERNIE-Bot", - Object: "model", - Created: 1677649963, - OwnedBy: "baidu", - Permission: permission, - Root: "ERNIE-Bot", - Parent: nil, - }, - { - Id: "ERNIE-Bot-turbo", - Object: "model", - Created: 1677649963, - OwnedBy: "baidu", - Permission: permission, - Root: "ERNIE-Bot-turbo", - Parent: nil, - }, - { - Id: "ERNIE-Bot-4", - Object: "model", - Created: 1677649963, - OwnedBy: "baidu", - Permission: permission, - Root: "ERNIE-Bot-4", - Parent: nil, - }, - { - Id: "Embedding-V1", - Object: "model", - Created: 1677649963, - OwnedBy: "baidu", - Permission: permission, - Root: "Embedding-V1", - Parent: nil, - }, - { - Id: "PaLM-2", - Object: "model", - Created: 1677649963, - OwnedBy: "google", - Permission: permission, - Root: "PaLM-2", - Parent: nil, - }, - { - Id: "gemini-pro", - Object: "model", - Created: 1677649963, - OwnedBy: "google", - Permission: permission, - Root: "gemini-pro", - Parent: nil, - }, - { - Id: "gemini-pro-vision", - Object: "model", - Created: 1677649963, - OwnedBy: "google", - Permission: permission, - Root: "gemini-pro-vision", - Parent: nil, - }, - { - Id: "chatglm_turbo", - Object: "model", - Created: 1677649963, - OwnedBy: "zhipu", - Permission: permission, - Root: "chatglm_turbo", - Parent: nil, - }, - { - Id: "chatglm_pro", - Object: "model", - Created: 1677649963, - OwnedBy: "zhipu", - Permission: permission, - Root: "chatglm_pro", - Parent: nil, - }, - { - Id: "chatglm_std", - Object: "model", - Created: 1677649963, - OwnedBy: "zhipu", - Permission: permission, - Root: "chatglm_std", - Parent: nil, - }, - { - Id: "chatglm_lite", - Object: "model", - Created: 1677649963, - OwnedBy: "zhipu", - Permission: permission, - Root: "chatglm_lite", - Parent: nil, - }, - { - Id: "qwen-turbo", - Object: "model", - Created: 1677649963, - OwnedBy: "ali", - Permission: permission, - Root: "qwen-turbo", - Parent: nil, - }, - { - Id: "qwen-plus", - Object: "model", - Created: 1677649963, - OwnedBy: "ali", - Permission: permission, - Root: "qwen-plus", - Parent: nil, - }, - { - Id: "qwen-max", - Object: "model", - Created: 1677649963, - OwnedBy: "ali", - Permission: permission, - Root: "qwen-max", - Parent: nil, - }, - { - Id: "qwen-max-longcontext", - Object: "model", - Created: 1677649963, - OwnedBy: "ali", - Permission: permission, - Root: "qwen-max-longcontext", - Parent: nil, - }, - { - Id: "text-embedding-v1", - Object: "model", - Created: 1677649963, - OwnedBy: "ali", - Permission: permission, - Root: "text-embedding-v1", - Parent: nil, - }, - { - Id: "SparkDesk", - Object: "model", - Created: 1677649963, - OwnedBy: "xunfei", - Permission: permission, - Root: "SparkDesk", - Parent: nil, - }, - { - Id: "360GPT_S2_V9", - Object: "model", - Created: 1677649963, - OwnedBy: "360", - Permission: permission, - Root: "360GPT_S2_V9", - Parent: nil, - }, - { - Id: "embedding-bert-512-v1", - Object: "model", - Created: 1677649963, - OwnedBy: "360", - Permission: permission, - Root: "embedding-bert-512-v1", - Parent: nil, - }, - { - Id: "embedding_s1_v1", - Object: "model", - Created: 1677649963, - OwnedBy: "360", - Permission: permission, - Root: "embedding_s1_v1", - Parent: nil, - }, - { - Id: "semantic_similarity_s1_v1", - Object: "model", - Created: 1677649963, - OwnedBy: "360", - Permission: permission, - Root: "semantic_similarity_s1_v1", - Parent: nil, - }, - { - Id: "hunyuan", - Object: "model", - Created: 1677649963, - OwnedBy: "tencent", - Permission: permission, - Root: "hunyuan", - Parent: nil, - }, + keys := make([]string, 0, len(common.ModelRatio)) + for k := range common.ModelRatio { + keys = append(keys, k) } + sort.Strings(keys) + + for _, modelId := range keys { + openAIModels = append(openAIModels, OpenAIModels{ + Id: modelId, + Object: "model", + Created: 1677649963, + OwnedBy: nil, + Permission: nil, + Root: nil, + Parent: nil, + }) + } + openAIModelsMap = make(map[string]OpenAIModels) for _, model := range openAIModels { openAIModelsMap[model.Id] = model @@ -585,6 +68,35 @@ func init() { } func ListModels(c *gin.Context) { + groupName := c.GetString("group") + + models, err := model.CacheGetGroupModels(groupName) + if err != nil { + common.AbortWithMessage(c, http.StatusServiceUnavailable, err.Error()) + return + } + sort.Strings(models) + + groupOpenAIModels := make([]OpenAIModels, 0, len(models)) + for _, modelId := range models { + groupOpenAIModels = append(groupOpenAIModels, OpenAIModels{ + Id: modelId, + Object: "model", + Created: 1677649963, + OwnedBy: nil, + Permission: nil, + Root: nil, + Parent: nil, + }) + } + + c.JSON(200, gin.H{ + "object": "list", + "data": groupOpenAIModels, + }) +} + +func ListModelsForAdmin(c *gin.Context) { c.JSON(200, gin.H{ "object": "list", "data": openAIModels, diff --git a/model/ability.go b/model/ability.go index 3da83be8..1a53b3ef 100644 --- a/model/ability.go +++ b/model/ability.go @@ -39,6 +39,22 @@ func GetRandomSatisfiedChannel(group string, model string) (*Channel, error) { return &channel, err } +func GetGroupModels(group string) ([]string, error) { + var models []string + groupCol := "`group`" + trueVal := "1" + if common.UsingPostgreSQL { + groupCol = `"group"` + trueVal = "true" + } + + err := DB.Model(&Ability{}).Where(groupCol+" = ? and enabled = ? ", group, trueVal).Distinct("model").Pluck("model", &models).Error + if err != nil { + return nil, err + } + return models, nil +} + func (channel *Channel) AddAbilities() error { models_ := strings.Split(channel.Models, ",") groups_ := strings.Split(channel.Group, ",") diff --git a/model/cache.go b/model/cache.go index c6d0c70a..e26995fe 100644 --- a/model/cache.go +++ b/model/cache.go @@ -213,3 +213,22 @@ func CacheGetRandomSatisfiedChannel(group string, model string) (*Channel, error idx := rand.Intn(endIdx) return channels[idx], nil } + +func CacheGetGroupModels(group string) ([]string, error) { + if !common.MemoryCacheEnabled { + return GetGroupModels(group) + } + channelSyncLock.RLock() + defer channelSyncLock.RUnlock() + + groupModels := group2model2channels[group] + if groupModels == nil { + return nil, errors.New("group not found") + } + + models := make([]string, 0) + for model := range groupModels { + models = append(models, model) + } + return models, nil +} diff --git a/router/api-router.go b/router/api-router.go index 162675ce..5cb2407e 100644 --- a/router/api-router.go +++ b/router/api-router.go @@ -67,7 +67,7 @@ func SetApiRouter(router *gin.Engine) { { channelRoute.GET("/", controller.GetAllChannels) channelRoute.GET("/search", controller.SearchChannels) - channelRoute.GET("/models", controller.ListModels) + channelRoute.GET("/models", controller.ListModelsForAdmin) channelRoute.GET("/:id", controller.GetChannel) channelRoute.GET("/test", controller.TestAllChannels) channelRoute.GET("/test/:id", controller.TestChannel) diff --git a/router/relay-router.go b/router/relay-router.go index 171864c6..519a764b 100644 --- a/router/relay-router.go +++ b/router/relay-router.go @@ -11,7 +11,7 @@ func SetRelayRouter(router *gin.Engine) { router.Use(middleware.CORS()) // https://platform.openai.com/docs/api-reference/introduction modelsRouter := router.Group("/v1/models") - modelsRouter.Use(middleware.TokenAuth()) + modelsRouter.Use(middleware.TokenAuth(), middleware.Distribute()) { modelsRouter.GET("", controller.ListModels) modelsRouter.GET("/:model", controller.RetrieveModel)