Update main.py

main.py

@@ -25,6 +25,9 @@ from fastapi.staticfiles import StaticFiles
 
 from bearer_token import BearerTokenGenerator
 
+from fastapi import Depends, HTTPException, Security
+from fastapi.security import HTTPAuthorizationCredentials, HTTPBearer
+
 # 模型列表
 MODELS = ["gpt-4o", "gpt-4o-mini", "claude-3-5-sonnet", "claude"]
 
@@ -121,6 +124,18 @@ def is_base64_image(url: str) -> bool:
     """
     return url.startswith("data:image/")
 
+# 添加 HTTPBearer 实例
+security = HTTPBearer()
+
+# 添加 API_KEY 验证函数
+def verify_api_key(credentials: HTTPAuthorizationCredentials = Security(security)):
+    api_key = os.environ.get("API_KEY")
+    if api_key is None:
+        raise HTTPException(status_code=500, detail="API_KEY not set in environment variables")
+    if credentials.credentials != api_key:
+        raise HTTPException(status_code=401, detail="Invalid API key")
+    return credentials.credentials
+
 # 根路径GET请求处理
 @app.get("/", response_class=HTMLResponse)
 async def read_root():
@@ -140,7 +155,7 @@ async def read_root():
 
 # 聊天完成处理
 @app.post("/ai/v1/chat/completions")
-async def chat_completions(request: Request, background_tasks: BackgroundTasks):
+async def chat_completions(request: Request, background_tasks: BackgroundTasks, api_key: str = Depends(verify_api_key)):
     """
     处理聊天完成请求
     """
@@ -380,7 +395,7 @@ async def chat_completions(request: Request, background_tasks: BackgroundTasks):
 
 # 图像生成处理
 @app.post("/ai/v1/images/generations")
-async def images_generations(request: Request):
+async def images_generations(request: Request, api_key: str = Depends(verify_api_key)):
     """
     处理图像生成请求
     """
@@ -544,10 +559,13 @@ def main():
     parser.add_argument('--base_url', type=str, default='http://localhost', help='Base URL for accessing images')
     parser.add_argument('--port', type=int, default=INITIAL_PORT, help='服务器监听端口')
     args = parser.parse_args()
-
     base_url = args.base_url
     port = args.port
 
+    # 检查 API_KEY 是否设置
+    if not os.environ.get("API_KEY"):
+        print("警告: API_KEY 环境变量未设置。客户端验证将无法正常工作。")
+
     # 确保 images 目录存在
     if not os.path.exists("images"):
         os.makedirs("images")