diff --git a/.env.example b/.env.example index 357eb69..49587d1 100644 --- a/.env.example +++ b/.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 diff --git a/Dockerfile b/Dockerfile index 1d8ad53..1cb49cc 100644 --- a/Dockerfile +++ b/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 diff --git a/README.md b/README.md index 94e64bc..ecd2095 100644 --- a/README.md +++ b/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` - 安全提示: 生产环境建议禁用 diff --git a/app/api/gemini_routes.py b/app/api/gemini_routes.py index 46df9d6..1596e18 100644 --- a/app/api/gemini_routes.py +++ b/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: - for name in settings.MODEL_SEARCH: + if settings.SEARCH_MODELS: + 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: - for name in settings.MODEL_IMAGE: + if settings.IMAGE_MODELS: + for name in settings.IMAGE_MODELS: model = model_mapping.get(name, None) if not model: continue diff --git a/app/api/openai_routes.py b/app/api/openai_routes.py index bf3d464..3e1f728 100644 --- a/app/api/openai_routes.py +++ b/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() diff --git a/app/core/config.py b/app/core/config.py index daa9cdf..bfa1835 100644 --- a/app/core/config.py +++ b/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"] - MODEL_IMAGE: List[str] = ["gemini-2.0-flash-exp"] + SEARCH_MODELS: 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 diff --git a/app/services/model_service.py b/app/services/model_service.py index 6dc7573..7cf97be 100644 --- a/app/services/model_service.py +++ b/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): - self.model_search = model_search - self.model_image = model_image + def __init__(self, search_models: list, image_models: list): + self.search_models = search_models + 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