Update app.py

app.py

@@ -2,11 +2,13 @@ import gradio as gr
 from transformers import AutoTokenizer, AutoModelForCausalLM
 import torch
 import json
-from fastapi import FastAPI, Request
+from fastapi import FastAPI, Request, HTTPException, Security, Depends
+from fastapi.security import APIKeyHeader
 from fastapi.responses import JSONResponse
 import logging
 import time
-import re
+import os
+from typing import Optional
 
 # 配置日志
 logging.basicConfig(level=logging.INFO)
@@ -15,10 +17,18 @@ logger = logging.getLogger(__name__)
 # 全局变量
 model = None
 tokenizer = None
-device = "cpu"  # 默认使用CPU
+device = "cpu"
+
+# 安全配置
+# 从环境变量读取配置，默认启用安全认证
+TEST_MODE: bool = os.getenv("TEST_MODE", "false").lower() == "true"
+API_KEYS = os.getenv("API_KEYS", "your-secret-key-1,your-secret-key-2").split(",")
+
+# 初始化API密钥头认证
+api_key_header = APIKeyHeader(name="X-API-Key", auto_error=False)
 
 def load_model():
-    """加载模型 - 不使用device_map以避免accelerate依赖"""
+    """加载模型"""
     global model, tokenizer, device
     
     if model is not None:
@@ -27,7 +37,6 @@ def load_model():
     try:
         model_name = "deepseek-ai/DeepSeek-R1-Distill-Qwen-1.5B"
         
-        # 检查是否有GPU可用
         if torch.cuda.is_available():
             device = "cuda"
             logger.info("检测到GPU可用，将使用GPU加速")
@@ -36,13 +45,11 @@ def load_model():
         
         tokenizer = AutoTokenizer.from_pretrained(model_name)
         
-        # 加载模型
         model = AutoModelForCausalLM.from_pretrained(
             model_name,
-            torch_dtype=torch.float32  # 使用float32确保兼容性
+            torch_dtype=torch.float32
         )
         
-        # 将模型移动到设备
         model = model.to(device)
         
         if tokenizer.pad_token is None:
@@ -54,8 +61,40 @@ def load_model():
         logger.error(f"模型加载失败: {e}")
         return False
 
+def verify_api_key(
+    request_key_header: Optional[str] = Security(api_key_header) if not TEST_MODE else None,
+) -> str:
+    """
+    API密钥验证依赖函数
+    支持测试模式和生产模式切换[1,3](@ref)
+    """
+    logger.info(f"当前安全模式: {'测试模式' if TEST_MODE else '生产模式'}")
+    
+    # 测试模式：跳过认证
+    if TEST_MODE:
+        logger.info("测试模式下跳过API密钥验证")
+        return "test_mode_bypass"
+    
+    # 生产模式：严格验证
+    if request_key_header is None:
+        logger.warning("请求头中缺少API密钥")
+        raise HTTPException(
+            status_code=401,
+            detail="缺少API密钥，请在请求头中添加 X-API-Key"
+        )
+    
+    if request_key_header not in API_KEYS:
+        logger.warning(f"无效的API密钥尝试: {request_key_header}")
+        raise HTTPException(
+            status_code=401,
+            detail="无效的API密钥"
+        )
+    
+    logger.info("API密钥验证通过")
+    return request_key_header
+
 def generate_response(message):
-    """生成模型响应 - 修正版"""
+    """生成模型响应"""
     if not load_model():
         return "模型加载失败，请稍后重试"
     
@@ -65,7 +104,6 @@ def generate_response(message):
             {"role": "user", "content": message}
         ]
         
-        # 使用tokenizer的apply_chat_template方法
         formatted_prompt = tokenizer.apply_chat_template(
             prompt,
             tokenize=False,
@@ -77,14 +115,14 @@ def generate_response(message):
             formatted_prompt, 
             return_tensors="pt", 
             truncation=True, 
-            max_length=512  # 减少输入长度
+            max_length=512
         ).to(device)
         
-        # 生成回复 - 减少生成长度
+        # 生成回复
         with torch.no_grad():
             outputs = model.generate(
                 **inputs,
-                max_new_tokens=128,  # 减少生成长度
+                max_new_tokens=128,
                 temperature=0.7,
                 top_p=0.9,
                 do_sample=True,
@@ -93,40 +131,43 @@ def generate_response(message):
                 repetition_penalty=1.1
             )
         
-        # 解码回复 - 跳过特殊标记
+        # 解码回复
         response = tokenizer.decode(
             outputs[0][inputs.input_ids.shape[-1]:], 
             skip_special_tokens=True
         )
         
-        # 关键修正：移除模型内部的思考过程
-        # 只保留最终回复内容
-        if "</think>" in response:
-            # 提取最终回复部分
-            final_response = response.split("</think>")[-1].strip()
-            # 移除可能的换行符和多余空格
-            final_response = re.sub(r'\n+', ' ', final_response).strip()
-            return final_response
-        else:
-            return response.strip()
+        return response.strip()
         
     except Exception as e:
         logger.error(f"生成回复时出错: {str(e)}")
         return f"生成回复时出错: {str(e)}"
 
 # 创建FastAPI应用
-app = FastAPI()
+app = FastAPI(title="DeepSeek-R1 API服务", description="带API密钥认证的大模型服务")
+
+# API健康检查端点（无需认证）
+@app.get("/")
+async def root():
+    return {"message": "DeepSeek-R1 API服务运行中", "timestamp": int(time.time())}
 
-# 添加API端点
+@app.get("/health")
+async def health_check():
+    return {"status": "healthy", "model_loaded": model is not None}
+
+# 受保护的聊天API端点
 @app.post("/api/chat")
-async def chat_api(request: Request):
-    """OpenAI兼容的聊天API端点"""
+async def chat_api(
+    request: Request,
+    api_key: str = Depends(verify_api_key)  # 添加API密钥依赖
+):
+    """OpenAI兼容的聊天API端点（需要API密钥认证）"""
     try:
         # 解析请求数据
         data = await request.json()
         messages = data.get("messages", [])
         model_name = data.get("model", "deepseek-r1")
-        max_tokens = data.get("max_tokens", 128)  # 默认128
+        max_tokens = data.get("max_tokens", 128)
         temperature = data.get("temperature", 0.7)
         
         # 提取最后一条用户消息
@@ -178,12 +219,15 @@ async def chat_api(request: Request):
 
 # 创建Gradio界面
 with gr.Blocks(title="DeepSeek-R1 API服务") as demo:
-    gr.Markdown("""
+    gr.Markdown(f"""
     # DeepSeek-R1 API 服务
     *基于DeepSeek-R1-Distill-Qwen-1.5B模型*
     
+    ## 安全状态: {'🔓 测试模式（认证已禁用）' if TEST_MODE else '🔒 生产模式（认证已启用）'}
+    
     ## API端点信息
-    - **OpenAI兼容端点**: `/api/chat`
+    - **聊天端点**: `/api/chat` (需要API密钥认证)
+    - **健康检查**: `/health` (公开)
     - **模型名称**: `deepseek-r1`
     """)
     
@@ -196,11 +240,9 @@ with gr.Blocks(title="DeepSeek-R1 API服务") as demo:
     
     # 处理函数
     def respond(message):
-        """处理用户输入"""
         if not message.strip():
             return ""
         
-        # 生成响应
         response = generate_response(message)
         return response