Update app.py

app.py

@@ -134,21 +134,33 @@ async def access_control(request: Request, call_next):
 
 # 流式响应生成器
 async def stream_generator(response):
-    async for line in response.aiter_lines():
-        if line.strip():  # 确保行不是空的
-            try:
-                # 如果是JSON格式的数据，直接转发
-                if line.startswith('{'):
-                    yield f"data: {line}\n\n"
-                # 如果已经是SSE格式，直接发送
-                elif line.startswith('data: '):
-                    yield f"{line}\n\n"
-                # 其他情况，添加data前缀
-                else:
-                    yield f"data: {line}\n\n"
-            except Exception as e:
-                print(f"Error processing line: {str(e)}")
-                continue
+    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"
+
 
 # 修改 handle_api_request 函数中的超时设置
 async def handle_api_request(url: str, headers: dict, method: str = "GET", body: dict = None, for_chat: bool = False):
@@ -312,37 +324,55 @@ async def add_key(request: Request):
 @app.post("/api/v1/chat/completions")
 async def chat_completions(request: Request):
     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",
+            "Accept": "text/event-stream" if body_json.get("stream") else "application/json"
         }
         
-        print("Starting chat completion request...")
+        # 构建目标URL
+        url = f"{Config.OPENAI_API_BASE}/chat/completions"
         
-        async with httpx.AsyncClient() as client:
-            key = get_chat_key()
-            headers["Authorization"] = f"Bearer {key}"
-            
-            async with client.stream(
-                "POST",
-                f"{Config.OPENAI_API_BASE}/chat/completions",
+        async with httpx.AsyncClient(timeout=60.0) as client:
+            response = await client.post(
+                url,
                 headers=headers,
-                json=body_json,
-                timeout=None
-            ) as response:
-                if response.status_code != 200:
-                    raise HTTPException(status_code=response.status_code, detail="API request failed")
-                
+                json=body_json
+            )
+            
+            # 检查响应状态
+            if response.status_code != 200:
+                return Response(
+                    content=response.text,
+                    status_code=response.status_code,
+                    media_type=response.headers.get("content-type", "application/json")
+                )
+            
+            # 处理流式响应
+            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"
+                    }
                 )
-        
+            
+            # 处理普通响应
+            return Response(
+                content=response.text,
+                media_type=response.headers.get("content-type", "application/json")
+            )
+                
     except Exception as e:
-        print(f"Chat completion error: {str(e)}")
+        print(f"Chat Error: {str(e)}")
         raise HTTPException(status_code=500, detail=str(e))
 
 # 代理其他请求