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

Fix api key risk
This commit is contained in:
Wenjie Zhang 2025-03-15 01:49:31 +08:00 committed by GitHub
commit 19453bb218
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
5 changed files with 186 additions and 70 deletions

View File

@ -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

View File

@ -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}")

View File

@ -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():

Binary file not shown.

Before

Width:  |  Height:  |  Size: 28 KiB

After

Width:  |  Height:  |  Size: 20 KiB

View File

@ -5,11 +5,8 @@
<p>配置文件也可以在 <code>saves/config/base.yaml</code> 中修改</p>
</template>
<template #actions>
<a-button type="primary" v-if="isNeedRestart" @click="sendRestart">
<ReloadOutlined />需要重启
</a-button>
<a-button v-else @click="sendRestart">
<ReloadOutlined />重启服务
<a-button :type="isNeedRestart ? 'primary' : 'default'" @click="sendRestart" :icon="h(ReloadOutlined)">
{{ isNeedRestart ? '需要重启' : '重启服务' }}
</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 的模型比如 vllmOllama</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>
<a :href="modelNames[item].url" target="_blank"><InfoCircleOutlined /></a>
</div>
<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-title-container">
<h3>{{ modelNames[item].name }}</h3>
<a :href="modelNames[item].url" target="_blank" class="model-url">
<InfoCircleOutlined />
</a>
</div>
<a-button
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>
<a :href="modelNames[item].url" target="_blank"><InfoCircleOutlined /></a>
<div class="model-icon">
<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) {
configStore.config.custom_models = []
}
let 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 (configStore.config.custom_models.find(item => item.custom_id == customModel.custom_id)) {
message.error('模型ID已存在')
if (customModel.edit_type === 'add') {
if (custom_models.find(item => item.custom_id === customModel.custom_id)) {
message.error('模型ID已存在')
return
}
const modelHash = generateRandomHash(4)
model_info.custom_id = `${customModel.name}-${modelHash}`
configStore.config.custom_models.push(model_info)
custom_models.push(model_info)
} else {
configStore.config.custom_models = configStore.config.custom_models.map(item => {
if (item.custom_id == customModel.custom_id) {
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
})
// custom_id
custom_models = custom_models.map(item => item.custom_id === customModel.custom_id ? model_info : item);
}
customModel.visible = false
await configStore.setConfigValue('custom_models', configStore.config.custom_models)
message.success('添加自定义模型成功')
customModel.visible = false;
await configStore.setConfigValue('custom_models', custom_models);
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;
height: 24px;
width: 28px;
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 {