Update main.py

main.py

@@ -1,14 +1,14 @@
 import os
 import json
-import time
+import asyncio
 import httpx
+import traceback
 from fastapi import FastAPI, Request
-from fastapi.responses import JSONResponse
+from fastapi.responses import StreamingResponse, JSONResponse
 from fastapi.middleware.cors import CORSMiddleware
 
 app = FastAPI()
 
-# 允许跨域请求
 app.add_middleware(
     CORSMiddleware,
     allow_origins=["*"],
@@ -16,7 +16,6 @@ app.add_middleware(
     allow_headers=["*"],
 )
 
-# 从环境变量获取你的配置
 AIROPS_API_KEY = os.environ.get("AIROPS_API_KEY", "")
 AIROPS_APP_ID = os.environ.get("AIROPS_APP_ID", "")
 INPUT_VAR = os.environ.get("INPUT_VAR_NAME", "input")
@@ -24,79 +23,82 @@ INPUT_VAR = os.environ.get("INPUT_VAR_NAME", "input")
 @app.post("/v1/chat/completions")
 async def chat_completions(request: Request):
     try:
-        body = await request.json()
-    except:
-        body = {}
-        
-    messages = body.get("messages", [])
-    
-    # 【酒馆特供版】将所有历史消息和设定拼接成一个长剧本
-    full_context = ""
-    for msg in messages:
-        role = msg.get("role", "user").upper()
-        content = msg.get("content", "")
+        try:
+            body = await request.json()
+        except:
+            body = {}
+            
+        messages = body.get("messages", [])
         
-        # 统一格式化标签，让 AI 知道是谁在说话
-        if role == "SYSTEM":
-            full_context += f"[系统设定及角色卡]:\n{content}\n\n"
-        elif role == "ASSISTANT":
-            full_context += f"[角色回复]:\n{content}\n\n"
-        else:
-            full_context += f"[我的输入]:\n{content}\n\n"
+        # 将所有历史消息和设定拼接成一个长剧本
+        full_context = ""
+        for msg in messages:
+            # 强制转换为字符串，防止特殊格式报错
+            role = str(msg.get("role", "user")).upper()
+            content = str(msg.get("content", ""))
             
-    # 如果没有收到任何消息，给个默认值
-    if not full_context.strip():
-        full_context = "[我的输入]:\n你好"
-    
-    headers = {
-        "Authorization": f"Bearer {AIROPS_API_KEY}",
-        "Content-Type": "application/json"
-    }
-    
-    # 构造 AirOps 需要的数据格式
-    payload = {
-        "inputs": {
-            INPUT_VAR: full_context.strip()
+            if role == "SYSTEM":
+                full_context += f"[系统设定及角色卡]:\n{content}\n\n"
+            elif role == "ASSISTANT":
+                full_context += f"[角色回复]:\n{content}\n\n"
+            else:
+                full_context += f"[我的输入]:\n{content}\n\n"
+                
+        if not full_context.strip():
+            full_context = "[我的输入]:\n你好"
+        
+        headers = {
+            "Authorization": f"Bearer {AIROPS_API_KEY}",
+            "Content-Type": "application/json"
         }
-    }
-    
-    # 设置 120 秒超时，防止长文本生成失败
-    async with httpx.AsyncClient(timeout=120.0) as client:
-        url = f"https://app.airops.com/public_api/airops_apps/{AIROPS_APP_ID}/execute"
         
-        try:
-            response = await client.post(url, headers=headers, json=payload)
-            if response.status_code == 200:
-                data = response.json()
-                # 提取 AirOps 返回的文本结果
-                result_text = data.get("result", "")
-                
-                if isinstance(result_text, (dict, list)):
-                    result_text = json.dumps(result_text, ensure_ascii=False)
+        payload = {
+            "inputs": {
+                INPUT_VAR: full_context.strip()
+            }
+        }
+        
+        async def generate():
+            # 同样保持 300 秒超时，给 AI 充分的思考时间
+            async with httpx.AsyncClient(timeout=300.0) as client:
+                url = f"https://app.airops.com/public_api/airops_apps/{AIROPS_APP_ID}/execute"
                 
-                # 核心修改：取消流式，直接返回标准的 OpenAI 格式 JSON 响应
-                return {
-                    "id": f"chatcmpl-airops-{int(time.time())}",
-                    "object": "chat.completion",
-                    "created": int(time.time()),
-                    "model": "airops-model",
-                    "choices": [
-                        {
-                            "index": 0,
-                            "message": {
-                                "role": "assistant",
-                                "content": result_text
-                            },
-                            "finish_reason": "stop"
-                        }
-                    ],
-                    "usage": {
-                        "prompt_tokens": 0,
-                        "completion_tokens": 0,
-                        "total_tokens": 0
-                    }
-                }
-            else:
-                return JSONResponse(status_code=500, content={"error": f"AirOps 接口报错: {response.text}"})
-        except Exception as e:
-            return JSONResponse(status_code=500, content={"error": f"代理服务器内部错误: {str(e)}"})
+                try:
+                    response = await client.post(url, headers=headers, json=payload)
+                    
+                    if response.status_code == 200:
+                        data = response.json()
+                        result_text = data.get("result", "")
+                        
+                        if isinstance(result_text, (dict, list)):
+                            result_text = json.dumps(result_text, ensure_ascii=False)
+                        
+                        # 模拟 OpenAI 流式打字机效果
+                        chunk_size = 3
+                        for i in range(0, len(result_text), chunk_size):
+                            chunk = result_text[i:i+chunk_size]
+                            openai_chunk = {
+                                "id": "chatcmpl-airops",
+                                "object": "chat.completion.chunk",
+                                "choices": [{"delta": {"content": chunk}}]
+                            }
+                            yield f"data: {json.dumps(openai_chunk)}\n\n"
+                            await asyncio.sleep(0.01)
+                            
+                        yield "data: [DONE]\n\n"
+                    else:
+                        print(f"AirOps Error: {response.status_code} - {response.text}")
+                        err = {"error": f"AirOps 接口报错: {response.status_code}"}
+                        yield f"data: {json.dumps(err)}\n\n"
+                        yield "data: [DONE]\n\n"
+                except Exception as e:
+                    traceback.print_exc()
+                    err = {"error": f"代理服务器内部错误: {repr(e)}"}
+                    yield f"data: {json.dumps(err)}\n\n"
+                    yield "data: [DONE]\n\n"
+
+        return StreamingResponse(generate(), media_type="text/event-stream")
+        
+    except Exception as e:
+        traceback.print_exc()
+        return JSONResponse(status_code=500, content={"error": f"请求解析错误: {repr(e)}"})