Update app.py

app.py

@@ -2,7 +2,8 @@ import os
 import time
 import json
 import logging
-from fastapi import FastAPI, Request, HTTPException
+from fastapi import FastAPI, Request, HTTPException, Depends, status
+from fastapi.security import HTTPBearer, HTTPAuthorizationCredentials
 from fastapi.responses import JSONResponse
 from transformers import AutoTokenizer, AutoModelForCausalLM
 import torch
@@ -21,6 +22,40 @@ MODEL_NAME = "Qwen/Qwen1.5-0.5B-Chat"
 MAX_TOKENS = 256
 DEVICE = "cpu"  # 强制使用CPU
 
+# API密钥配置
+# 从环境变量获取API密钥，如果没有设置则使用默认值
+API_KEYS = os.getenv("API_KEYS", "your-secret-key-1,your-secret-key-2").split(",")
+# 是否启用API密钥验证
+API_AUTH_ENABLED = os.getenv("API_AUTH_ENABLED", "true").lower() == "true"
+
+# 创建Bearer认证方案
+security = HTTPBearer()
+
+def verify_api_key(credentials: HTTPAuthorizationCredentials = Depends(security)):
+    """验证API密钥"""
+    # 如果未启用认证，则跳过验证
+    if not API_AUTH_ENABLED:
+        return True
+    
+    # 检查Bearer令牌格式
+    if not credentials.scheme == "Bearer":
+        raise HTTPException(
+            status_code=status.HTTP_401_UNAUTHORIZED,
+            detail="Invalid authentication scheme. Use 'Bearer' token",
+            headers={"WWW-Authenticate": "Bearer"},
+        )
+    
+    # 检查API密钥是否有效
+    api_key = credentials.credentials
+    if api_key not in API_KEYS:
+        raise HTTPException(
+            status_code=status.HTTP_401_UNAUTHORIZED,
+            detail="Invalid API key",
+            headers={"WWW-Authenticate": "Bearer"},
+        )
+    
+    return True
+
 def load_model():
     """极简模型加载"""
     global model, tokenizer
@@ -81,6 +116,14 @@ def generate_response(messages):
             add_generation_prompt=True
         )
         
+        # 确保text是字符串
+        if not isinstance(text, str):
+            # 如果返回的是列表，则连接成字符串
+            if isinstance(text, list):
+                text = "".join(text)
+            else:
+                text = str(text)
+        
         # 编码输入
         inputs = tokenizer(text, return_tensors="pt", truncation=True, max_length=1024)
         inputs = inputs.to(DEVICE)
@@ -112,25 +155,46 @@ def generate_response(messages):
         return {"error": f"生成失败: {str(e)}"}
 
 # 创建极简FastAPI应用
-app = FastAPI(title="Qwen1.5-0.5B API", version="1.0")
+app = FastAPI(
+    title="Qwen1.5-0.5B API",
+    version="1.0",
+    description="带有API密钥验证的Qwen1.5-0.5B-Chat API服务"
+)
 
 # 启动时加载模型
 @app.on_event("startup")
 async def startup_event():
     load_model()
+    logger.info(f"API认证状态: {'已启用' if API_AUTH_ENABLED else '已禁用'}")
+    if API_AUTH_ENABLED:
+        logger.info(f"有效的API密钥数量: {len(API_KEYS)}")
 
-# 健康检查端点（OpenClaw可能需要）
+# 健康检查端点（无需认证）
 @app.get("/health")
 async def health_check():
     return {
         "status": "healthy" if model is not None else "loading",
         "model_loaded": model is not None,
+        "api_auth_enabled": API_AUTH_ENABLED,
         "timestamp": int(time.time())
     }
 
-# OpenAI兼容的聊天端点
+# 根端点（无需认证）
+@app.get("/")
+async def root():
+    return {
+        "message": "Qwen1.5-0.5B-Chat API服务运行中",
+        "model_loaded": model is not None,
+        "api_auth_enabled": API_AUTH_ENABLED,
+        "endpoint": "/v1/chat/completions"
+    }
+
+# OpenAI兼容的聊天端点（需要认证）
 @app.post("/v1/chat/completions")
-async def chat_completion(request: Request):
+async def chat_completion(
+    request: Request,
+    auth_valid: bool = Depends(verify_api_key)
+):
     """极简版OpenAI兼容端点"""
     try:
         # 解析请求
@@ -185,13 +249,14 @@ async def chat_completion(request: Request):
             }
         )
 
-# 根端点
-@app.get("/")
-async def root():
+# 添加一个简单的测试端点（需要认证）
+@app.post("/v1/test")
+async def test_endpoint(auth_valid: bool = Depends(verify_api_key)):
+    """测试端点，验证API密钥是否有效"""
     return {
-        "message": "Qwen1.5-0.5B-Chat API服务运行中",
-        "model_loaded": model is not None,
-        "endpoint": "/v1/chat/completions"
+        "status": "success",
+        "message": "API密钥验证通过",
+        "timestamp": int(time.time())
     }
 
 if __name__ == "__main__":