添加获取安全API的方法

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