Update app.py

app.py

@@ -2,34 +2,25 @@ import os
 import time
 import json
 import logging
-import gc
-from typing import List, Dict, Any, Union
-
 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
+import gc
 
-# -----------------------------------------------------------------------------
-# 日志配置
-# -----------------------------------------------------------------------------
-logging.basicConfig(
-    level=logging.INFO,
-    format='%(asctime)s - %(levelname)s - %(message)s'
-)
+# 极简日志配置
+logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
 logger = logging.getLogger(__name__)
 
-# -----------------------------------------------------------------------------
-# 全局变量与配置
-# -----------------------------------------------------------------------------
+# 全局变量
 model = None
 tokenizer = None
 
 # 配置
 MODEL_NAME = "Qwen/Qwen1.5-0.5B-Chat"
 MAX_TOKENS = 256
-DEVICE = "cpu"  # 强制使用CPU，如需GPU请改为 "cuda"
+DEVICE = "cpu"  # 强制使用CPU
 
 # API密钥配置
 API_KEYS = os.getenv("API_KEYS", "your-secret-key-1,your-secret-key-2").split(",")
@@ -38,10 +29,6 @@ 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:
@@ -68,9 +55,6 @@ def load_model():
     """极简模型加载"""
     global model, tokenizer
     
-    if model is not None:
-        return True
-
     try:
         logger.info(f"开始加载模型: {MODEL_NAME}")
         
@@ -93,7 +77,7 @@ def load_model():
             trust_remote_code=True
         )
         
-        # 移动到设备
+        # 移动到CPU
         model = model.to(DEVICE)
         model.eval()  # 设置为评估模式
         
@@ -104,67 +88,85 @@ def load_model():
         logger.error(f"模型加载失败: {e}")
         return False
 
-def generate_chat_response(messages: List[Dict[str, str]], max_tokens=256, temperature=0.7):
+def apply_chat_template(messages):
     """
-    使用 apply_chat_template 生成响应，支持完整的对话历史。
-    这是 Hugging Face 推荐的标准方式。
+    把 OpenAI 格式的 messages 转为 Qwen 的 chat template 格式
     """
+    text = ""
+    for msg in messages:
+        role = msg.get("role", "").lower()
+        content = msg.get("content", "").strip()
+        if not content:
+            continue
+            
+        if role == "system":
+            text += f"<|im_start|>system\n{content}<|im_end|>\n"
+        elif role == "user":
+            text += f"<|im_start|>user\n{content}<|im_end|>\n"
+        elif role == "assistant":
+            text += f"<|im_start|>assistant\n{content}<|im_end|>\n"
+        else:
+            # 忽略其他 role
+            continue
+    
+    # 最后加上 assistant 的开头
+    text += "<|im_start|>assistant\n"
+    return text
+
+
+def generate_chat_response(messages, max_tokens=256, temperature=0.7):
+    """生成完整对话回复"""
     if model is None or tokenizer is None:
         return {"error": "模型未加载"}
-    
+
     try:
-        # 使用 apply_chat_template 构建输入，它会自动处理 <|im_start|> 等格式
-        # tokenize=False 返回字符串，方便调试和确保输入类型正确
-        text = tokenizer.apply_chat_template(
-            messages,
-            tokenize=False,
-            add_generation_prompt=True
-        )
-        
-        # 记录生成的文本提示（用于调试）
-        # logger.info(f"生成的Prompt片段: {text[:100]}...")
+        # 转换为 Qwen 的对话格式
+        prompt = apply_chat_template(messages)
         
-        # 编码输入 - 直接传入字符串，不放入列表，避免某些 tokenizer 版本的批处理歧义
-        model_inputs = tokenizer(
-            text, 
-            return_tensors="pt"
-        ).to(DEVICE)
-        
-        # 生成响应
+        logger.info(f"输入文本类型: {type(prompt)}, 长度: {len(prompt)}")
+        logger.debug(f"完整prompt前100字符: {prompt[:100]}...")
+
+        # 分词（注意这里使用列表包一层字符串）
+        inputs = tokenizer(
+            [prompt],               # 必须是 list[str]
+            return_tensors="pt",
+            truncation=True,
+            max_length=3072,        # Qwen1.5 支持较长上下文
+            padding=True
+        )
+        inputs = {k: v.to(DEVICE) for k, v in inputs.items()}
+
+        # 生成
         with torch.no_grad():
-            generated_ids = model.generate(
-                model_inputs.input_ids,
+            outputs = model.generate(
+                **inputs,
                 max_new_tokens=min(max_tokens, MAX_TOKENS),
                 do_sample=True,
                 temperature=temperature,
-                top_p=0.9,
-                pad_token_id=tokenizer.eos_token_id
+                top_p=0.85,
+                repetition_penalty=1.05,     # 轻微防止重复
+                pad_token_id=tokenizer.eos_token_id,
+                eos_token_id=tokenizer.eos_token_id,
             )
-        
-        # 获取新生成的token（去掉输入的token）
-        generated_ids = [
-            output_ids[len(input_ids):] 
-            for input_ids, output_ids in zip(model_inputs.input_ids, generated_ids)
-        ]
 
-        # 解码响应
-        response_text = tokenizer.batch_decode(generated_ids, skip_special_tokens=True)[0]
-        
-        # 清理内存
-        del model_inputs, generated_ids
-        if torch.cuda.is_available():
-            torch.cuda.empty_cache()
+        # 只取新生成的 token
+        generated_ids = outputs[0][inputs["input_ids"].shape[1]:]
+        response = tokenizer.decode(generated_ids, skip_special_tokens=True)
+
+        # 清理可能的结束标记
+        response = response.split("<|im_end|>")[0].strip()
+
+        # 内存清理
+        del inputs, outputs
         gc.collect()
-        
-        return {"text": response_text}
-        
+
+        return {"text": response}
+
     except Exception as e:
-        logger.error(f"生成响应失败: {e}", exc_info=True)
-        return {"error": f"生成失败: {str(e)}"}
+        logger.error(f"生成失败: {str(e)}", exc_info=True)
+        return {"error": str(e)}
 
-# -----------------------------------------------------------------------------
-# FastAPI 应用
-# -----------------------------------------------------------------------------
+# 创建极简FastAPI应用
 app = FastAPI(
     title="OpenAI API兼容服务",
     version="1.0",
@@ -179,7 +181,7 @@ async def startup_event():
     if API_AUTH_ENABLED:
         logger.info(f"有效的API密钥数量: {len(API_KEYS)}")
 
-# 健康检查端点
+# 健康检查端点（无需认证）
 @app.get("/health")
 async def health_check():
     return {
@@ -189,6 +191,31 @@ async def health_check():
         "timestamp": int(time.time())
     }
 
+# 根端点（无需认证）
+@app.get("/")
+async def root():
+    return {
+        "message": "OpenAI API兼容服务运行中",
+        "model_loaded": model is not None,
+        "api_auth_enabled": API_AUTH_ENABLED,
+        "endpoints": {
+            "v1": "/v1",
+            "chat_completions": "/v1/chat/completions"
+        }
+    }
+
+# 添加/v1端点（OpenClaw可能需要）
+@app.get("/v1")
+async def v1_root():
+    return {
+        "message": "OpenAI v1 API端点",
+        "endpoints": {
+            "models": "/v1/models",
+            "chat_completions": "/v1/chat/completions"
+        }
+    }
+
+# 添加模型列表端点（OpenAI兼容）
 @app.get("/v1/models")
 async def list_models():
     """返回可用的模型列表"""
@@ -204,42 +231,37 @@ async def list_models():
         ]
     }
 
-# OpenAI Chat Completions端点
+# OpenAI Chat Completions端点（主要端点）
 @app.post("/v1/chat/completions")
 async def create_chat_completion(
     request: Request,
     auth_valid: bool = Depends(verify_api_key)
 ):
-    """OpenAI Chat Completions API兼容端点"""
     try:
-        # 解析请求
         data = await request.json()
         messages = data.get("messages", [])
-        model_name = data.get("model", "qwen1.5-0.5b-chat")
+        model_name = data.get("model", MODEL_NAME)
         max_tokens = data.get("max_tokens", MAX_TOKENS)
         temperature = data.get("temperature", 0.7)
-        
-        logger.info(f"收到Chat Completions请求: model={model_name}, messages_count={len(messages)}")
-        
-        # 检查消息格式
+
+        logger.info(f"收到请求: model={model_name}, messages_count={len(messages)}")
+
         if not messages or not isinstance(messages, list):
-            return JSONResponse(
-                status_code=400,
-                content={"error": {"message": "无效的消息格式", "type": "invalid_request_error"}}
-            )
-        
-        # 使用新的生成函数，直接传递 messages 列表
+            raise ValueError("messages 必须是非空列表")
+
+        # 生成回复
         result = generate_chat_response(messages, max_tokens, temperature)
-        
+
         if "error" in result:
-            return JSONResponse(
-                status_code=500,
-                content={"error": {"message": result["error"], "type": "internal_error"}}
-            )
-        
-        # 返回OpenAI Chat Completions兼容格式
+            raise RuntimeError(result["error"])
+
+        # 计算粗略 token 数（仅供参考）
+        prompt_text = "".join([m["content"] for m in messages if m.get("content")])
+        prompt_tokens = len(tokenizer.encode(prompt_text)) if tokenizer else 0
+        completion_tokens = len(tokenizer.encode(result["text"])) if tokenizer else 0
+
         response_data = {
-            "id": f"chatcmpl-{int(time.time())}",
+            "id": f"chatcmpl-{int(time.time()*1000)}",
             "object": "chat.completion",
             "created": int(time.time()),
             "model": model_name,
@@ -254,33 +276,109 @@ async def create_chat_completion(
                 }
             ],
             "usage": {
-                "prompt_tokens": -1, 
-                "completion_tokens": -1,
-                "total_tokens": -1
+                "prompt_tokens": prompt_tokens,
+                "completion_tokens": completion_tokens,
+                "total_tokens": prompt_tokens + completion_tokens
             }
         }
-        
-        logger.info(f"成功生成响应: {len(result['text'])} 字符")
+
         return response_data
-        
+
     except Exception as e:
-        logger.error(f"Chat Completions API错误: {e}", exc_info=True)
+        logger.error(f"Chat Completions 错误: {str(e)}", exc_info=True)
         return JSONResponse(
             status_code=500,
             content={
                 "error": {
-                    "message": f"内部服务器错误: {str(e)}",
-                    "type": "internal_error"
+                    "message": str(e),
+                    "type": "internal_server_error"
+                }
+            }
+        )
+
+# 添加兼容性端点（为不同版本的OpenClaw提供支持）
+@app.post("/chat/completions")
+async def legacy_chat_completion(
+    request: Request,
+    auth_valid: bool = Depends(verify_api_key)
+):
+    """兼容旧版本OpenClaw的端点"""
+    # 直接转发到/v1/chat/completions
+    return await create_chat_completion(request, auth_valid)
+
+# 添加通用聊天端点
+@app.post("/api/chat")
+async def generic_chat_api(
+    request: Request,
+    auth_valid: bool = Depends(verify_api_key)
+):
+    """通用聊天API端点"""
+    try:
+        # 解析请求
+        data = await request.json()
+        messages = data.get("messages", [])
+        
+        # 检查消息格式
+        if not messages or not isinstance(messages, list):
+            return JSONResponse(
+                status_code=400,
+                content={
+                    "error": "无效的消息格式"
+                }
+            )
+        
+        # 提取用户消息
+        user_message = ""
+        for msg in messages:
+            if isinstance(msg, dict) and msg.get("role") == "user":
+                user_message = msg.get("content", "")
+                break
+        
+        if not user_message:
+            return JSONResponse(
+                status_code=400,
+                content={
+                    "error": "未找到用户消息"
+                }
+            )
+        
+        # 生成响应
+        result = generate_completion(user_message)
+        
+        if "error" in result:
+            return JSONResponse(
+                status_code=500,
+                content={
+                    "error": result["error"]
                 }
+            )
+        
+        # 返回通用格式
+        return {
+            "choices": [{
+                "message": {
+                    "content": result["text"]
+                }
+            }]
+        }
+        
+    except Exception as e:
+        logger.error(f"通用聊天API错误: {e}")
+        return JSONResponse(
+            status_code=500,
+            content={
+                "error": f"内部服务器错误: {str(e)}"
             }
         )
 
 if __name__ == "__main__":
     import uvicorn
+    
+    # 极简UVicorn配置
     uvicorn.run(
         app, 
         host="0.0.0.0", 
         port=7860,
-        workers=1,
+        workers=1,  # 单worker减少内存占用
         log_level="info"
     )