chore: 规范变量名

.env.example

@@ -1,9 +1,10 @@
 API_KEYS=["AIzaSyxxxxxxxxxxxxxxxxxxx","AIzaSyxxxxxxxxxxxxxxxxxxx"]
 ALLOWED_TOKENS=["sk-123456"]
 # AUTH_TOKEN=sk-123456
+TEST_MODEL=gemini-1.5-flash
+IMAGE_MODELS=["gemini-2.0-flash-exp"]
+SEARCH_MODELS=["gemini-2.0-flash-exp","gemini-2.0-pro-exp"]
 FILTERED_MODELS=["gemini-1.0-pro-vision-latest", "gemini-pro-vision", "chat-bison-001", "text-bison-001", "embedding-gecko-001"]
-MODEL_SEARCH=["gemini-2.0-flash-exp","gemini-2.0-pro-exp"]
-MODEL_IMAGE=["gemini-2.0-flash-exp"]
 TOOLS_CODE_EXECUTION_ENABLED=false
 SHOW_SEARCH_LINK=true
 SHOW_THINKING_PROCESS=true

Dockerfile

@@ -11,7 +11,7 @@ ENV API_KEYS='["your_api_key_1"]'
 ENV ALLOWED_TOKENS='["your_token_1"]'
 ENV BASE_URL=https://generativelanguage.googleapis.com/v1beta
 ENV TOOLS_CODE_EXECUTION_ENABLED=false
-ENV MODEL_SEARCH='["gemini-2.0-flash-exp"]'
+ENV IMAGE_MODELS='["gemini-2.0-flash-exp"]'
 
 # Expose port
 EXPOSE 8000

README.md

@@ -64,10 +64,13 @@
     AUTH_TOKEN=""  # 超级管理员token，具有所有权限，默认使用 ALLOWED_TOKENS 的第一个
 
     # 模型功能配置
-    MODEL_SEARCH=["gemini-2.0-flash-exp"]  # 支持搜索功能的模型列表
+    TEST_MODEL="gemini-1.5-flash" # 用于测试密钥是否可用的模型名
+    SEARCH_MODELS=["gemini-2.0-flash-exp"]  # 支持搜索功能的模型列表
+    IMAGE_MODELS=["gemini-2.0-flash-exp"] # 支持绘图功能的模型列表
     TOOLS_CODE_EXECUTION_ENABLED=false  # 是否启用代码执行工具，默认false
     SHOW_SEARCH_LINK=true  # 是否在响应中显示搜索结果链接，默认true
     SHOW_THINKING_PROCESS=true  # 是否显示模型思考过程，默认true
+    FILTERED_MODELS=["gemini-1.0-pro-vision-latest", "gemini-pro-vision", "chat-bison-001", "text-bison-001", "embedding-gecko-001"] # 被禁用的模型列表
 
     # 图片生成配置
     PAID_KEY="your-paid-api-key"  # 付费版API Key，用于图片生成等高级功能
@@ -115,9 +118,17 @@
 
    #### 模型功能配置
 
-    - `MODEL_SEARCH`: 搜索功能支持的模型
+    - `TEST_MODEL`: 用于测试密钥可用性的模型
+      - 默认值: `"gemini-1.5-flash"`
+    - `SEARCH_MODELS`: 搜索功能支持的模型
       - 默认值: `["gemini-2.0-flash-exp"]`
       - 说明: 仅列表中的模型可使用搜索功能
+    - `IMAGE_MODELS`: 绘图功能支持的模型
+      - 默认值: `["gemini-2.0-flash-exp"]`
+      - 说明: 仅列表中的模型可使用绘图功能
+    - `FILTERED_MODELS`: 被禁用的模型列表
+      - 默认值: `["gemini-1.0-pro-vision-latest", "gemini-pro-vision", "chat-bison-001", "text-bison-001", "embedding-gecko-001"]`
+      - 说明: 列表中的模型将被禁用
     - `TOOLS_CODE_EXECUTION_ENABLED`: 代码执行功能
       - 默认值: `false`
       - 安全提示: 生产环境建议禁用

app/api/gemini_routes.py

@@ -23,7 +23,7 @@ async def get_key_manager():
 async def get_next_working_key_wrapper(key_manager: KeyManager = Depends(get_key_manager)):
     return await key_manager.get_next_working_key()
 
-model_service = ModelService(settings.MODEL_SEARCH,settings.MODEL_IMAGE)
+model_service = ModelService(settings.SEARCH_MODELS,settings.IMAGE_MODELS)
 
 
 @router.get("/models")
@@ -41,8 +41,8 @@ async def list_models(_=Depends(security_service.verify_key),
     model_mapping = {x.get("name", "").split("/", maxsplit=1)[1]: x for x in models_json["models"]}
 
     # 添加搜索模型
-    if settings.MODEL_SEARCH:
+    if settings.SEARCH_MODELS:
-        for name in settings.MODEL_SEARCH:
+        for name in settings.SEARCH_MODELS:
             model = model_mapping.get(name, None)
             if not model:
                 continue
@@ -56,8 +56,8 @@ async def list_models(_=Depends(security_service.verify_key),
             models_json["models"].append(item)
 
     # 添加图像生成模型
-    if settings.MODEL_IMAGE:
+    if settings.IMAGE_MODELS:
-        for name in settings.MODEL_IMAGE:
+        for name in settings.IMAGE_MODELS:
             model = model_mapping.get(name, None)
             if not model:
                 continue

app/api/openai_routes.py

@@ -17,7 +17,7 @@ logger = get_openai_logger()
 
 # 初始化服务
 security_service = SecurityService(settings.ALLOWED_TOKENS, settings.AUTH_TOKEN)
-model_service = ModelService(settings.MODEL_SEARCH,settings.MODEL_IMAGE)
+model_service = ModelService(settings.SEARCH_MODELS,settings.IMAGE_MODELS)
 embedding_service = EmbeddingService(settings.BASE_URL)
 image_create_service = ImageCreateService()
 

app/core/config.py

@@ -6,8 +6,8 @@ class Settings(BaseSettings):
     API_KEYS: List[str]
     ALLOWED_TOKENS: List[str]
     BASE_URL: str = "https://generativelanguage.googleapis.com/v1beta"
-    MODEL_SEARCH: List[str] = ["gemini-2.0-flash-exp"]
+    SEARCH_MODELS: List[str] = ["gemini-2.0-flash-exp"]
-    MODEL_IMAGE: List[str] = ["gemini-2.0-flash-exp"]
+    IMAGE_MODELS: List[str] = ["gemini-2.0-flash-exp"]
     FILTERED_MODELS: List[str] = ["gemini-1.0-pro-vision-latest", "gemini-pro-vision", "chat-bison-001", "text-bison-001", "embedding-gecko-001"]
     TOOLS_CODE_EXECUTION_ENABLED: bool = False
     SHOW_SEARCH_LINK: bool = True

app/services/model_service.py

@@ -7,9 +7,9 @@ from app.core.config import settings
 logger = get_model_logger()
 
 class ModelService:
-    def __init__(self, model_search: list, model_image: list):
+    def __init__(self, search_models: list, image_models: list):
-        self.model_search = model_search
+        self.search_models = search_models
-        self.model_image = model_image
+        self.image_models = image_models
         self.base_url = "https://generativelanguage.googleapis.com/v1beta"
         self.filtered_models = settings.FILTERED_MODELS
 
@@ -65,11 +65,11 @@ class ModelService:
             }
             openai_format["data"].append(openai_model)
 
-            if model_id in self.model_search:
+            if model_id in self.search_models:
                 search_model = openai_model.copy()
                 search_model["id"] = f"{model_id}-search"
                 openai_format["data"].append(search_model)
-            if model_id in self.model_image:
+            if model_id in self.image_models:
                 image_model = openai_model.copy()
                 image_model["id"] = f"{model_id}-image"
                 openai_format["data"].append(image_model)
@@ -87,9 +87,9 @@ class ModelService:
         model = model.strip()
         if model.endswith("-search"):
             model = model[:-7]
-            return model in settings.MODEL_SEARCH
+            return model in self.search_models
         if model.endswith("-image"):
             model = model[:-6]
-            return model in settings.MODEL_IMAGE
+            return model in self.image_models
 
-        return True
+        return model not in self.filtered_models