feat: 填入相关模型

Qianxia666 · May 12, 2024 · 2dbf50d · 2dbf50d
1 parent d8c0060
commit 2dbf50d
Show file tree

Hide file tree

Showing 11 changed files with 111 additions and 115 deletions.
diff --git a/common/constants.go b/common/constants.go
@@ -208,6 +208,8 @@ const (
 	ChannelTypeLingYiWanWu    = 31
 	ChannelTypeAws            = 33
 	ChannelTypeCohere         = 34
+
+	ChannelTypeDummy // this one is only for count, do not add any channel after this
 )
 
 var ChannelBaseURLs = []string{

diff --git a/controller/model.go b/controller/model.go
@@ -3,14 +3,17 @@ package controller
 import (
 	"fmt"
 	"github.com/gin-gonic/gin"
+	"log"
 	"net/http"
+	"one-api/common"
 	"one-api/constant"
 	"one-api/dto"
 	"one-api/model"
 	"one-api/relay"
 	"one-api/relay/channel/ai360"
-	"one-api/relay/channel/moonshot"
 	"one-api/relay/channel/lingyiwanwu"
+	"one-api/relay/channel/moonshot"
+	relaycommon "one-api/relay/common"
 	relayconstant "one-api/relay/constant"
 )
 
@@ -43,6 +46,7 @@ type OpenAIModels struct {
 
 var openAIModels []OpenAIModels
 var openAIModelsMap map[string]OpenAIModels
+var channelId2Models map[int][]string
 
 func init() {
 	var permission []OpenAIModelPermission
@@ -85,7 +89,7 @@ func init() {
 			Id:         modelName,
 			Object:     "model",
 			Created:    1626777600,
-			OwnedBy:    "360",
+			OwnedBy:    ai360.ChannelName,
 			Permission: permission,
 			Root:       modelName,
 			Parent:     nil,
@@ -128,6 +132,18 @@ func init() {
 	for _, model := range openAIModels {
 		openAIModelsMap[model.Id] = model
 	}
+	channelId2Models = make(map[int][]string)
+	for i := 1; i <= common.ChannelTypeDummy; i++ {
+		apiType := relayconstant.ChannelType2APIType(i)
+		if apiType == -1 || apiType == relayconstant.APITypeAIProxyLibrary {
+			continue
+		}
+		log.Println(apiType)
+		meta := &relaycommon.RelayInfo{ChannelType: i}
+		adaptor := relay.GetAdaptor(apiType)
+		adaptor.Init(meta, dto.GeneralOpenAIRequest{})
+		channelId2Models[i] = adaptor.GetModelList()
+	}
 }
 
 func ListModels(c *gin.Context) {
@@ -148,15 +164,22 @@ func ListModels(c *gin.Context) {
 		}
 	}
 	c.JSON(200, gin.H{
-		"object": "list",
-		"data":   userOpenAiModels,
+		"success": true,
+		"data":    userOpenAiModels,
 	})
 }
 
 func ChannelListModels(c *gin.Context) {
 	c.JSON(200, gin.H{
-		"object": "list",
-		"data":   openAIModels,
+		"success": true,
+		"data":    openAIModels,
+	})
+}
+
+func DashboardListModels(c *gin.Context) {
+	c.JSON(200, gin.H{
+		"success": true,
+		"data":    channelId2Models,
 	})
 }
 

diff --git a/relay/channel/ai360/constants.go b/relay/channel/ai360/constants.go
@@ -6,3 +6,5 @@ var ModelList = []string{
 	"embedding_s1_v1",
 	"semantic_similarity_s1_v1",
 }
+
+var ChannelName = "ai360"
diff --git a/relay/channel/ollama/constants.go b/relay/channel/ollama/constants.go
@@ -1,5 +1,7 @@
 package ollama
 
-var ModelList []string
+var ModelList = []string{
+	"llama3-7b",
+}
 
 var ChannelName = "ollama"
diff --git a/relay/channel/openai/constant.go b/relay/channel/openai/constant.go
@@ -6,7 +6,7 @@ var ModelList = []string{
 	"gpt-3.5-turbo-instruct",
 	"gpt-4", "gpt-4-0314", "gpt-4-0613", "gpt-4-1106-preview", "gpt-4-0125-preview",
 	"gpt-4-32k", "gpt-4-32k-0314", "gpt-4-32k-0613",
-	"gpt-4-turbo-preview",
+	"gpt-4-turbo-preview", "gpt-4-turbo", "gpt-4-turbo-2024-04-09",
 	"gpt-4-vision-preview",
 	"text-embedding-ada-002", "text-embedding-3-small", "text-embedding-3-large",
 	"text-curie-001", "text-babbage-001", "text-ada-001", "text-davinci-002", "text-davinci-003",

diff --git a/relay/constant/api_type.go b/relay/constant/api_type.go
@@ -25,8 +25,18 @@ const (
 )
 
 func ChannelType2APIType(channelType int) int {
-	apiType := APITypeOpenAI
+	apiType := -1
 	switch channelType {
+	case common.ChannelTypeOpenAI:
+		apiType = APITypeOpenAI
+	case common.ChannelTypeAzure:
+		apiType = APITypeOpenAI
+	case common.ChannelTypeMoonshot:
+		apiType = APITypeOpenAI
+	case common.ChannelTypeLingYiWanWu:
+		apiType = APITypeOpenAI
+	case common.ChannelType360:
+		apiType = APITypeOpenAI
 	case common.ChannelTypeAnthropic:
 		apiType = APITypeAnthropic
 	case common.ChannelTypeBaidu:

diff --git a/router/api-router.go b/router/api-router.go
@@ -14,6 +14,7 @@ func SetApiRouter(router *gin.Engine) {
 	apiRouter.Use(middleware.GlobalAPIRateLimit())
 	{
 		apiRouter.GET("/status", controller.GetStatus)
+		apiRouter.GET("/models", middleware.UserAuth(), controller.DashboardListModels)
 		apiRouter.GET("/status/test", middleware.AdminAuth(), controller.TestStatus)
 		apiRouter.GET("/notice", controller.GetNotice)
 		apiRouter.GET("/about", controller.GetAbout)

diff --git a/web/src/components/ChannelsTable.js b/web/src/components/ChannelsTable.js
@@ -31,6 +31,7 @@ import {
 } from '@douyinfe/semi-ui';
 import EditChannel from '../pages/Channel/EditChannel';
 import { IconTreeTriangleDown } from '@douyinfe/semi-icons';
+import { loadChannelModels } from './utils.js';
 
 function renderTimestamp(timestamp) {
   return <>{timestamp2string(timestamp)}</>;
@@ -354,27 +355,29 @@ const ChannelsTable = () => {
   };
 
   const copySelectedChannel = async (id) => {
-    const channelToCopy = channels.find(channel => String(channel.id) === String(id));
-    console.log(channelToCopy)
+    const channelToCopy = channels.find(
+      (channel) => String(channel.id) === String(id),
+    );
+    console.log(channelToCopy);
     channelToCopy.name += '_复制';
     channelToCopy.created_time = null;
     channelToCopy.balance = 0;
     channelToCopy.used_quota = 0;
     if (!channelToCopy) {
-        showError("渠道未找到，请刷新页面后重试。");
-        return;
+      showError('渠道未找到，请刷新页面后重试。');
+      return;
     }
     try {
-        const newChannel = {...channelToCopy, id: undefined};
-        const response = await API.post('/api/channel/', newChannel);
-        if (response.data.success) {
-            showSuccess("渠道复制成功");
-            await refresh();
-        } else {
-            showError(response.data.message);
-        }
+      const newChannel = { ...channelToCopy, id: undefined };
+      const response = await API.post('/api/channel/', newChannel);
+      if (response.data.success) {
+        showSuccess('渠道复制成功');
+        await refresh();
+      } else {
+        showError(response.data.message);
+      }
     } catch (error) {
-        showError("渠道复制失败: " + error.message);
+      showError('渠道复制失败: ' + error.message);
     }
   };
 
@@ -395,6 +398,7 @@ const ChannelsTable = () => {
         showError(reason);
       });
     fetchGroups().then();
+    loadChannelModels().then();
   }, []);
 
   const manageChannel = async (id, action, record, value) => {

diff --git a/web/src/components/utils.js b/web/src/components/utils.js
@@ -18,3 +18,32 @@ export async function onGitHubOAuthClicked(github_client_id) {
     `https://github.com/login/oauth/authorize?client_id=${github_client_id}&state=${state}&scope=user:email`,
   );
 }
+
+let channelModels = undefined;
+export async function loadChannelModels() {
+  const res = await API.get('/api/models');
+  const { success, data } = res.data;
+  if (!success) {
+    return;
+  }
+  channelModels = data;
+  localStorage.setItem('channel_models', JSON.stringify(data));
+}
+
+export function getChannelModels(type) {
+  if (channelModels !== undefined && type in channelModels) {
+    if (!channelModels[type]) {
+      return [];
+    }
+    return channelModels[type];
+  }
+  let models = localStorage.getItem('channel_models');
+  if (!models) {
+    return [];
+  }
+  channelModels = JSON.parse(models);
+  if (type in channelModels) {
+    return channelModels[type];
+  }
+  return [];
+}
diff --git a/web/src/constants/channel.constants.js b/web/src/constants/channel.constants.js
@@ -86,13 +86,13 @@ export const CHANNEL_OPTIONS = [
     label: '智谱 ChatGLM',
   },
   {
-    key: 16,
+    key: 26,
     text: '智谱 GLM-4V',
     value: 26,
     color: 'purple',
     label: '智谱 GLM-4V',
   },
-  { key: 16, text: 'Moonshot', value: 25, color: 'green', label: 'Moonshot' },
+  { key: 25, text: 'Moonshot', value: 25, color: 'green', label: 'Moonshot' },
   { key: 19, text: '360 智脑', value: 19, color: 'blue', label: '360 智脑' },
   { key: 23, text: '腾讯混元', value: 23, color: 'teal', label: '腾讯混元' },
   { key: 31, text: '零一万物', value: 31, color: 'green', label: '零一万物' },