Update app.py

app.py

@@ -8,6 +8,7 @@ from typing import List, Optional
 import requests
 from itertools import cycle
 import asyncio
+import time
 
 # 创建FastAPI应用
 app = FastAPI()
@@ -26,6 +27,7 @@ class Config:
     OPENAI_API_BASE = "https://api.x.ai/v1"
     KEYS_URL = os.getenv("KEYS_URL", "")
     WHITELIST_IPS = os.getenv("WHITELIST_IPS", "").split(",")
+    CHUNK_SIZE = 1  # 控制流式传输的块大小
 
 # 全局变量
 keys = []
@@ -85,37 +87,60 @@ async def list_models():
 
 # 流式响应生成器
 async def stream_generator(response):
-    async for chunk in response.aiter_bytes():
-        yield chunk
+    buffer = ""
+    try:
+        async for chunk in response.aiter_bytes():
+            chunk_str = chunk.decode('utf-8')
+            buffer += chunk_str
+            
+            # 处理buffer中的完整事件
+            while '\n\n' in buffer:
+                event, buffer = buffer.split('\n\n', 1)
+                if event.startswith('data: '):
+                    data = event[6:]  # 移除 "data: " 前缀
+                    if data.strip() == '[DONE]':
+                        yield f"data: [DONE]\n\n"
+                    else:
+                        try:
+                            # 解析JSON并重新格式化
+                            json_data = json.loads(data)
+                            yield f"data: {json.dumps(json_data)}\n\n"
+                            # 添加小延迟使流更平滑
+                            await asyncio.sleep(0.01)
+                        except json.JSONDecodeError:
+                            print(f"JSON decode error for data: {data}")
+                            continue
+    except Exception as e:
+        print(f"Stream Error: {str(e)}")
+        yield f"data: {json.dumps({'error': str(e)})}\n\n"
 
 # 聊天完成路由
 @app.post("/api/v1/chat/completions")
 async def chat_completions(request: Request):
-    # 获取请求体
-    body = await request.body()
-    body_json = json.loads(body)
-    
-    # 获取headers
-    headers = {
-        "Authorization": f"Bearer {get_next_key()}",
-        "Content-Type": "application/json",
-        "Accept": "text/event-stream" if body_json.get("stream") else "application/json"
-    }
-    
-    # 构建目标URL
-    url = f"{Config.OPENAI_API_BASE}/chat/completions"
-    
-    print(f"Chat request to: {url}")
-    print(f"Headers: {headers}")
-    print(f"Body: {json.dumps(body_json)}")
-    
-    async with httpx.AsyncClient() as client:
-        try:
+    try:
+        # 获取请求体
+        body = await request.body()
+        body_json = json.loads(body)
+        
+        # 获取headers
+        headers = {
+            "Authorization": f"Bearer {get_next_key()}",
+            "Content-Type": "application/json",
+            "Accept": "text/event-stream" if body_json.get("stream") else "application/json"
+        }
+        
+        # 构建目标URL
+        url = f"{Config.OPENAI_API_BASE}/chat/completions"
+        
+        print(f"Chat request to: {url}")
+        print(f"Headers: {headers}")
+        print(f"Body: {json.dumps(body_json)}")
+        
+        async with httpx.AsyncClient(timeout=60.0) as client:
             response = await client.post(
                 url,
                 headers=headers,
-                json=body_json,
-                timeout=60.0
+                json=body_json
             )
             
             # 检查响应状态
@@ -131,7 +156,12 @@ async def chat_completions(request: Request):
             if body_json.get("stream"):
                 return StreamingResponse(
                     stream_generator(response),
-                    media_type="text/event-stream"
+                    media_type="text/event-stream",
+                    headers={
+                        "Cache-Control": "no-cache",
+                        "Connection": "keep-alive",
+                        "Content-Type": "text/event-stream"
+                    }
                 )
             
             # 处理普通响应
@@ -140,9 +170,9 @@ async def chat_completions(request: Request):
                 media_type=response.headers.get("content-type", "application/json")
             )
                 
-        except Exception as e:
-            print(f"Chat Error: {str(e)}")
-            raise HTTPException(status_code=500, detail=str(e))
+    except Exception as e:
+        print(f"Chat Error: {str(e)}")
+        raise HTTPException(status_code=500, detail=str(e))
 
 # 代理其他请求到X.AI
 @app.api_route("/api/v1/{path:path}", methods=["GET", "POST", "PUT", "DELETE"])