Merge pull request #76 from xerrors/fix-api-key-risk

Fix api key risk

src/config/__init__.py

@@ -11,6 +11,7 @@ MODEL_NAMES = _models["MODEL_NAMES"]
 EMBED_MODEL_INFO = _models["EMBED_MODEL_INFO"]
 RERANKER_LIST = _models["RERANKER_LIST"]
 
+DEFAULT_MOCK_API = 'this_is_mock_api_key_in_frontend'
 
 class SimpleConfig(dict):
 
@@ -174,3 +175,34 @@ class Config(SimpleConfig):
                 json.dump(self, f, indent=4)
 
         logger.info(f"Config file {self.filename} saved")
+
+    def get_safe_config(self):
+        """
+        获取安全的配置，即过滤掉 api_key
+        """
+        
+        config = json.loads(str(self))
+        
+        # 过滤掉 api_key
+        for model in config.get("custom_models", []):
+            model["api_key"] = DEFAULT_MOCK_API if model.get("api_key") else ""
+
+        return config
+    
+    def compare_custom_models(self, value):
+        """
+        比较 custom_models 中的 api_key，如果输入的 api_key 与当前的 api_key 相同，则不修改
+        如果输入的 api_key 为 DEFAULT_MOCK_API，则使用当前的 api_key
+        """
+        current_models_dict = {model["custom_id"]: model.get("api_key") for model in self.custom_models}
+
+        for i, model in enumerate(value):
+            input_custom_id = model.get("custom_id")
+            input_api_key = model.get("api_key")
+
+            if input_custom_id in current_models_dict:
+                current_api_key = current_models_dict[input_custom_id]
+                if input_api_key == DEFAULT_MOCK_API or input_api_key == current_api_key:
+                    value[i]["api_key"] = current_api_key
+
+        return value

src/models/chat_model.py

@@ -58,7 +58,7 @@ class OpenModel(OpenAIBase):
 class CustomModel(OpenAIBase):
     def __init__(self, model_info):
         model_name = model_info["name"]
-        api_key = model_info["api_key"] or "custom_model"
+        api_key = model_info.get("api_key", "custom_model")
         base_url = get_docker_safe_url(model_info["api_base"])
         logger.info(f"> Custom model: {model_name}, base_url: {base_url}")
 

src/routers/base_router.py

@@ -1,14 +1,10 @@
+from fastapi import Request, Body
 from fastapi import APIRouter
 
 base = APIRouter()
 
-from fastapi import FastAPI, HTTPException
-from fastapi.responses import JSONResponse
-from fastapi import Request, Body
-
-from src.core import HistoryManager
 from src.core.startup import startup
-from src.utils import logger
+from src.utils import hashstr, logger
 
 
 @base.get("/")
@@ -17,13 +13,16 @@ async def route_index():
 
 @base.get("/config")
 def get_config():
-    return startup.config
+    return startup.config.get_safe_config()
 
 @base.post("/config")
 async def update_config(key = Body(...), value = Body(...)):
+    if key == "custom_models":
+        value = startup.config.compare_custom_models(value)
+
     startup.config[key] = value
     startup.config.save()
-    return startup.config
+    return startup.config.get_safe_config()
 
 @base.post("/restart")
 async def restart():

web/src/views/SettingView.vue

@@ -5,11 +5,8 @@
         <p>配置文件也可以在 <code>saves/config/base.yaml</code> 中修改</p>
       </template>
       <template #actions>
-        <a-button type="primary" v-if="isNeedRestart" @click="sendRestart">
+        <a-button :type="isNeedRestart ? 'primary' : 'default'" @click="sendRestart" :icon="h(ReloadOutlined)">
-          <ReloadOutlined />需要重启
+          {{ isNeedRestart ? '需要重启' : '重启服务' }}
-        </a-button>
-        <a-button v-else @click="sendRestart">
-          <ReloadOutlined />重启服务
         </a-button>
       </template>
     </HeaderComponent>
@@ -121,12 +118,12 @@
                   >
                     <a-button type="text" :disabled="configStore.config.model_name == item.name"  @click.stop><DeleteOutlined /></a-button>
                   </a-popconfirm>
-                  <a-button type="text" @click.stop="handleEditCustomModel(item)"><EditOutlined /></a-button>
+                  <a-button type="text" @click.stop="prepareToEditCustomModel(item)"><EditOutlined /></a-button>
                 </div>
               </div>
               <div class="api_base">{{ item.api_base }}</div>
             </div>
-            <div class="card-models custom-model" @click="customModel.visible=true">
+            <div class="card-models custom-model" @click="prepareToAddCustomModel">
               <div class="card-models__header">
                 <div class="name"> + 添加模型</div>
               </div>
@@ -135,7 +132,7 @@
                 class="custom-model-modal"
                 v-model:open="customModel.visible"
                 :title="customModel.modelTitle"
-                @ok="handleAddCustomModel"
+                @ok="handleAddOrEditCustomModel"
                 @cancel="handleCancelCustomModel"
                 :okText="'确认'"
                 :cancelText="'取消'"
@@ -145,13 +142,13 @@
                 <p>添加的模型是兼容 OpenAI 的模型，比如 vllm，Ollama。</p>
                 <a-form :model="customModel" layout="vertical" >
                   <a-form-item label="模型名称" name="name" :rules="[{ required: true, message: '请输入模型名称' }]">
-                    <a-input v-model:value="customModel.name" />
+                    <a-input v-model:value="customModel.name" :disabled="customModel.edit_type == 'edit'"/>
                   </a-form-item>
                   <a-form-item label="API Base" name="api_base" :rules="[{ required: true, message: '请输入API Base' }]">
-                    <a-input v-model:value="customModel.api_base"/>
+                    <a-input v-model:value="customModel.api_base" />
                   </a-form-item>
                   <a-form-item label="API KEY" name="api_key">
-                    <a-input v-model:value="customModel.api_key" autocomplete="off"/>
+                    <a-input-password v-model:value="customModel.api_key" :visibilityToggle="false" autocomplete="new-password"/>
                   </a-form-item>
                 </a-form>
               </a-modal>
@@ -164,25 +161,47 @@
             <div :class="{'model-icon': true, 'available': modelStatus[item]}">
               <img :src="modelIcons[item]" alt="模型图标">
             </div>
-            <h3>{{ modelNames[item].name }}</h3>
+            <div class="model-title-container">
-            <a :href="modelNames[item].url" target="_blank"><InfoCircleOutlined /></a>
+              <h3>{{ modelNames[item].name }}</h3>
-          </div>
+              <a :href="modelNames[item].url" target="_blank" class="model-url">
-          <div class="card-body" v-if="modelStatus[item]">
+                <InfoCircleOutlined />
-            <div
+              </a>
-              :class="{'model_selected': modelProvider == item && configStore.config.model_name == model, 'card-models': true}"
+            </div>
-              v-for="(model, idx) in modelNames[item].models" :key="idx"
+            <a-button 
-              @click="handleChange('model_provider', item); handleChange('model_name', model)"
+              type="text" 
+              class="expand-button" 
+              @click="expandedModels[item] = !expandedModels[item]"
             >
-              <div class="model_name">{{ model }}</div>
+              <span class="icon-wrapper" :class="{'rotated': expandedModels[item]}">
+                <DownOutlined />
+              </span>
+            </a-button>
+          </div>
+          <div class="card-body-wrapper" :class="{'expanded': expandedModels[item]}">
+            <div class="card-body" v-if="modelStatus[item]">
+              <div
+                :class="{'model_selected': modelProvider == item && configStore.config.model_name == model, 'card-models': true}"
+                v-for="(model, idx) in modelNames[item].models" :key="idx"
+                @click="handleChange('model_provider', item); handleChange('model_name', model)"
+              >
+                <div class="model_name">{{ model }}</div>
+              </div>
             </div>
           </div>
         </div>
         <div class="model-provider-card" v-for="(item, key) in notModelKeys" :key="key">
           <div class="card-header">
-            <h3 style="font-weight: 400">{{ modelNames[item].name }}</h3>
+            <div class="model-icon">
-            <a :href="modelNames[item].url" target="_blank"><InfoCircleOutlined /></a>
+              <img :src="modelIcons[item]" alt="模型图标">
+            </div>
+            <div class="model-title-container">
+              <h3 style="font-weight: 400">{{ modelNames[item].name }}</h3>
+              <a :href="modelNames[item].url" target="_blank" class="model-url">
+                <InfoCircleOutlined />
+              </a>
+            </div>
             <div class="missing-keys">
-              <small>需配置</small> <span v-for="(key, idx) in modelNames[item].env" :key="idx">{{ key }}</span>
+              需配置<span v-for="(key, idx) in modelNames[item].env" :key="idx">{{ key }}</span>
             </div>
           </div>
         </div>
@@ -212,6 +231,8 @@ import {
   DeleteOutlined,
   EditOutlined,
   InfoCircleOutlined,
+  DownOutlined,
+  UpOutlined,
 } from '@ant-design/icons-vue';
 import HeaderComponent from '@/components/HeaderComponent.vue';
 import TableConfigComponent from '@/components/TableConfigComponent.vue';
@@ -228,7 +249,7 @@ const isNeedRestart = ref(false)
 const customModel = reactive({
   modelTitle: '添加自定义模型',
   visible: false,
-  custom_id: '', // Add this line
+  custom_id: '',
   name: '',
   api_key: '',
   api_base: '',
@@ -249,6 +270,12 @@ const notModelKeys = computed(() => {
   return Object.keys(modelStatus.value || {}).filter(key => !modelStatus.value?.[key])
 })
 
+// 模型展开状态管理，默认展开
+const expandedModels = reactive(modelKeys.value.reduce((acc, key) => {
+  acc[key] = true
+  return acc
+}, {}))
+
 const generateRandomHash = (length) => {
   let chars = 'abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789';
   let hash = '';
@@ -298,49 +325,35 @@ const handleChange = (key, e) => {
   configStore.setConfigValue(key, e)
 }
 
-const handleAddCustomModel = async () => {
+const handleAddOrEditCustomModel = async () => {
   if (!customModel.name || !customModel.api_base) {
-    message.error('请填写完整模型信息')
+    message.error('请填写完整的模型名称和API Base信息。')
     return
   }
 
-  if (!configStore.config.custom_models) {
+  let custom_models = configStore.config.custom_models || [];
-    configStore.config.custom_models = []
-  }
-
 
   const model_info = {
-    custom_id: customModel.custom_id,
+    custom_id: customModel.custom_id || `${customModel.name}-${generateRandomHash(4)}`,
     name: customModel.name,
     api_key: customModel.api_key,
     api_base: customModel.api_base,
   }
 
-  if (customModel.edit_type == 'add') {
+  if (customModel.edit_type === 'add') {
-  if (configStore.config.custom_models.find(item => item.custom_id == customModel.custom_id)) {
+    if (custom_models.find(item => item.custom_id === customModel.custom_id)) {
-    message.error('模型ID已存在')
+      message.error('模型ID已存在')
       return
     }
-    const modelHash = generateRandomHash(4)
+    custom_models.push(model_info)
-    model_info.custom_id = `${customModel.name}-${modelHash}`
-    configStore.config.custom_models.push(model_info)
   } else {
-    configStore.config.custom_models = configStore.config.custom_models.map(item => {
+    // 如果 custom_id 相同，则更新
-      if (item.custom_id == customModel.custom_id) {
+    custom_models = custom_models.map(item => item.custom_id === customModel.custom_id ? model_info : item);
-        if (item.name != customModel.name) {
-          const modelHash = generateRandomHash(4)
-          model_info.custom_id = `${customModel.name}-${modelHash}`
-          configStore.setConfigValue('model_name', model_info.custom_id)
-        }
-        return model_info
-      }
-      return item
-    })
   }
 
-  customModel.visible = false
+  customModel.visible = false;
-  await configStore.setConfigValue('custom_models', configStore.config.custom_models)
+  await configStore.setConfigValue('custom_models', custom_models);
-  message.success('添加自定义模型成功')
+  message.success(customModel.edit_type === 'add' ? '模型添加成功' : '模型修改成功');
 }
 
 const handleDeleteCustomModel = (custom_id) => {
@@ -348,21 +361,32 @@ const handleDeleteCustomModel = (custom_id) => {
   configStore.setConfigValue('custom_models', updatedModels);
 }
 
-const handleEditCustomModel = (item) => {
+const prepareToEditCustomModel = (item) => {
   customModel.modelTitle = '编辑自定义模型'
   customModel.custom_id = item.custom_id
+  customModel.visible = true
+  customModel.edit_type = 'edit'
   customModel.name = item.name
   customModel.api_key = item.api_key
   customModel.api_base = item.api_base
-  customModel.visible = true
-  customModel.edit_type = 'edit'
 }
 
-const handleCancelCustomModel = () => {
+const prepareToAddCustomModel = () => {
+  customModel.modelTitle = '添加自定义模型'
+  customModel.edit_type = 'add'
+  customModel.visible = true
+  clearCustomModel()
+}
+
+const clearCustomModel = () => {
   customModel.custom_id = ''
   customModel.name = ''
   customModel.api_key = ''
   customModel.api_base = ''
+}
+
+const handleCancelCustomModel = () => {
+  clearCustomModel()
   customModel.visible = false
 }
 
@@ -499,14 +523,31 @@ const sendRestart = () => {
       align-items: center;
       gap: 10px;
 
+      .model-title-container {
+        display: flex;
+        flex-direction: row;
+        align-items: center;
+        gap: 10px;
+        flex: 1;
+      }
+
+      .model-url {
+        font-size: 12px;
+        width: fit-content;
+        color: var(--gray-500);
+      }
+
       .model-icon {
-        width: 24px;
+        width: 28px;
-        height: 24px;
+        height: 28px;
+        
         // 灰度
         filter: grayscale(100%);
         img {
           width: 100%;
           height: 100%;
+          border-radius: 4px;
+          border: 1px solid var(--gray-300);
         }
 
         &.available {
@@ -514,10 +555,9 @@ const sendRestart = () => {
         }
       }
 
-
       h3 {
         margin: 0;
-        font-size: 1rem;
+        font-size: 0.9rem;
         font-weight: bold;
       }
 
@@ -546,12 +586,53 @@ const sendRestart = () => {
       }
 
       .missing-keys {
+        margin-top: 4px;
         color: var(--gray-600);
+        font-size: 12px;
         & > span {
-          margin-left: 10px;
+          margin-left: 6px;
           user-select: all;
+          background-color: var(--gray-100);
+          padding: 2px 6px;
+          border-radius: 4px;
+          color: var(--gray-800);
         }
       }
+
+      .expand-button {
+        margin-left: auto;
+        height: 32px;
+        width: 32px;
+        display: flex;
+        align-items: center;
+        justify-content: center;
+        padding: 0;
+        cursor: pointer;
+        color: var(--gray-500);
+
+        &:hover {
+          background-color: var(--gray-100);
+        }
+
+        .icon-wrapper {
+          display: inline-flex;
+          transition: transform 0.2s ease;
+          
+          &.rotated {
+            transform: rotate(180deg);
+          }
+        }
+      }
+    }
+
+    .card-body-wrapper {
+      max-height: 0;
+      overflow: hidden;
+      transition: max-height 0.2s ease-out;  // 先快后慢
+      
+      &.expanded {
+        max-height: 700px; /* 设置一个足够大的值 */
+      }
     }
 
     .card-body {
@@ -606,7 +687,7 @@ const sendRestart = () => {
           gap: 10px;
           .card-models__header {
             width: 100%;
-            height: 32px;
+            height: 24px;
             display: flex;
             justify-content: flex-start;
             align-items: center;
@@ -631,6 +712,10 @@ const sendRestart = () => {
           .api_base {
             font-size: 12px;
             color: var(--gray-600);
+            overflow: hidden;
+            text-overflow: ellipsis;
+            white-space: nowrap;
+            width: 100%;
           }
 
           &:hover {