Update app.py

app.py

@@ -10,98 +10,35 @@ import traceback
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 
-# ===== 测试代码开始 =====
-print("="*50)
-print("🔍 开始测试模型调用")
-print("="*50)
+# 初始化 FastAPI
+app = FastAPI()
 
+# 获取 token
 HF_TOKEN = os.getenv("HF_TOKEN")
 if not HF_TOKEN:
-    print("❌ 错误: HF_TOKEN 环境变量未设置!")
-else:
-    print(f"✅ HF_TOKEN 已设置 (长度: {len(HF_TOKEN)})")
-    if HF_TOKEN.startswith("hf_"):
-        print("✅ HF_TOKEN 格式正确")
-    else:
-        print("⚠️ 警告: HF_TOKEN 格式可能不正确，应以 hf_ 开头")
-
-# 测试你的模型
-model_id = "jiang1002/chatglm-6b-adgen"
-print(f"\n📊 正在测试模型 '{model_id}'...")
-
-# 方法1：测试 Hugging Face 免费推理（不指定 provider）
-try:
-    print("\n🔄 测试1: 使用 Hugging Face 免费推理...")
-    client1 = InferenceClient(token=HF_TOKEN)
-    response1 = client1.text_generation(
-        "你好",
-        model=model_id,
-        max_new_tokens=20
-    )
-    print(f"✅ 免费推理成功! 响应: {response1[:50]}...")
-except Exception as e:
-    print(f"❌ 免费推理失败: {str(e)}")
+    logger.error("❌ HF_TOKEN 环境变量未设置!")
 
-# 方法2：测试 auto provider（自动选择）
-try:
-    print("\n🔄 测试2: 使用 auto provider...")
-    client2 = InferenceClient(provider="auto", token=HF_TOKEN)
-    response2 = client2.chat.completions.create(
-        model=model_id,
-        messages=[{"role": "user", "content": "你好"}],
-        max_tokens=20
-    )
-    print(f"✅ auto provider 成功! 响应: {response2.choices[0].message.content[:50]}...")
-except Exception as e:
-    print(f"❌ auto provider 失败: {str(e)}")
+# 选择一个兼容的模型（改这里！）
+MODEL_ID = "meta-llama/Llama-3.2-3B-Instruct"  # 轻量级，免费，支持中文
+# MODEL_ID = "microsoft/Phi-3.5-mini-instruct"  # 备选1
+# MODEL_ID = "Qwen/Qwen2.5-7B-Instruct"  # 备选2
 
-# 方法3：测试 Groq（如果配置了）
-try:
-    print("\n🔄 测试3: 使用 Groq...")
-    client3 = InferenceClient(provider="groq", token=HF_TOKEN)
-    response3 = client3.chat.completions.create(
-        model=model_id,
-        messages=[{"role": "user", "content": "你好"}],
-        max_tokens=20
-    )
-    print(f"✅ Groq 成功! 响应: {response3.choices[0].message.content[:50]}...")
-except Exception as e:
-    print(f"❌ Groq 失败: {str(e)}")
+# 创建客户端（用 auto 模式）
+client = InferenceClient(provider="auto", token=HF_TOKEN)
 
-# 方法4：测试 Together AI（如果配置了）
+# 测试模型是否可用（启动时验证）
 try:
-    print("\n🔄 测试4: 使用 Together AI...")
-    client4 = InferenceClient(provider="together-ai", token=HF_TOKEN)
-    response4 = client4.chat.completions.create(
-        model=model_id,
+    logger.info(f"正在测试模型 {MODEL_ID}...")
+    test_response = client.chat.completions.create(
+        model=MODEL_ID,
         messages=[{"role": "user", "content": "你好"}],
-        max_tokens=20
+        max_tokens=10
     )
-    print(f"✅ Together AI 成功! 响应: {response4.choices[0].message.content[:50]}...")
+    logger.info(f"✅ 模型可用！响应: {test_response.choices[0].message.content}")
 except Exception as e:
-    print(f"❌ Together AI 失败: {str(e)}")
-
-print("\n" + "="*50)
-print("🔍 测试结束，继续启动应用...")
-print("="*50)
-# ===== 测试代码结束 =====
-
-# 初始化 FastAPI
-app = FastAPI()
-
-# 从环境变量获取 Hugging Face Token
-HF_TOKEN = os.getenv("HF_TOKEN")
-if not HF_TOKEN:
-    logger.warning("⚠️ 未设置 HF_TOKEN 环境变量，API 调用可能失败")
-
-# 初始化 InferenceClient
-# 这里用 provider="auto" 让系统自动选择可用提供商
-client = InferenceClient(provider="auto", token=HF_TOKEN)
+    logger.error(f"❌ 模型测试失败: {str(e)}")
 
-# 你的模型名称
-MODEL_ID = "jiang1002/chatglm-6b-adgen"  # 或者换成其他公开模型
-
-# --- 1. API 接口 ---
+# --- API 接口 ---
 @app.post("/generate")
 async def generate(request: Request):
     try:
@@ -109,7 +46,6 @@ async def generate(request: Request):
         prompt = data.get("text", "")
         messages = data.get("messages", [])
         
-        # 如果提供了完整的 messages 格式，就用它
         if messages:
             response = client.chat.completions.create(
                 model=MODEL_ID,
@@ -117,25 +53,22 @@ async def generate(request: Request):
             )
             result = response.choices[0].message.content
         else:
-            # 否则用简单的 prompt 格式
-            response = client.text_generation(
-                prompt,
+            response = client.chat.completions.create(
                 model=MODEL_ID,
-                max_new_tokens=512,
-                temperature=0.7
+                messages=[{"role": "user", "content": prompt}]
             )
-            result = response
+            result = response.choices[0].message.content
             
         return {"success": True, "result": result}
     except Exception as e:
         logger.error(f"API 调用失败: {str(e)}")
         return {"success": False, "error": str(e)}
 
-# --- 2. Gradio 聊天界面 ---
+# --- Gradio 聊天界面 ---
 def chat_func(message, history):
     """Gradio 聊天函数"""
     try:
-        # 将历史记录转换为 messages 格式
+        # 转换历史记录
         messages = []
         for human, assistant in history:
             messages.append({"role": "user", "content": human})
@@ -153,20 +86,20 @@ def chat_func(message, history):
         return response.choices[0].message.content
     except Exception as e:
         logger.error(f"聊天失败: {str(e)}")
-        logger.error(f"详细错误: {traceback.format_exc()}")  # 添加这行
-        return f"调用失败: {str(e)}\n\n{traceback.format_exc()}"
+        logger.error(traceback.format_exc())
+        return f"调用失败: {str(e)}"
 
 # 创建 Gradio 界面
 demo = gr.ChatInterface(
     fn=chat_func,
-    title="ChatGLM 广告生成助手 (使用 Inference Providers)",
-    description="后台使用 Hugging Face Inference Providers，无需本地 GPU"
+    title="AI 聊天助手 (使用 Inference Providers)",
+    description=f"后台模型: {MODEL_ID}"
 )
 
 # 挂载 Gradio
 app = gr.mount_gradio_app(app, demo, path="/")
 
-# 添加健康检查端点
+# 健康检查
 @app.get("/health")
 async def health():
     return {"status": "ok", "model": MODEL_ID}