Update app.py

app.py

@@ -1,11 +1,13 @@
 from fastapi import FastAPI, HTTPException, Request, Response
 from fastapi.middleware.cors import CORSMiddleware
+from fastapi.responses import StreamingResponse
 import httpx
 import os
 import json
 from typing import List, Optional
 import requests
 from itertools import cycle
+import asyncio
 
 # 创建FastAPI应用
 app = FastAPI()
@@ -81,9 +83,73 @@ async def list_models():
             print(f"Models Error: {str(e)}")
             raise HTTPException(status_code=500, detail=str(e))
 
+# 流式响应生成器
+async def stream_generator(response):
+    async for chunk in response.aiter_bytes():
+        yield chunk
+
+# 聊天完成路由
+@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:
+            response = await client.post(
+                url,
+                headers=headers,
+                json=body_json,
+                timeout=60.0
+            )
+            
+            # 检查响应状态
+            if response.status_code != 200:
+                print(f"Error response: {response.status_code} - {response.text}")
+                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"
+                )
+            
+            # 处理普通响应
+            return Response(
+                content=response.text,
+                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))
+
 # 代理其他请求到X.AI
 @app.api_route("/api/v1/{path:path}", methods=["GET", "POST", "PUT", "DELETE"])
 async def proxy(path: str, request: Request):
+    if path == "chat/completions":
+        return await chat_completions(request)
+        
     # 获取请求体
     body = await request.body()
     body_str = body.decode() if body else ""
@@ -100,11 +166,6 @@ async def proxy(path: str, request: Request):
     # 构建目标URL
     url = f"{Config.OPENAI_API_BASE}/{path}"
     
-    print(f"Proxying request to: {url}")
-    print(f"Method: {request.method}")
-    print(f"Headers: {headers}")
-    print(f"Body: {body_str}")
-    
     async with httpx.AsyncClient() as client:
         try:
             response = await client.request(
@@ -115,25 +176,11 @@ async def proxy(path: str, request: Request):
                 content=body_str if body_str else None
             )
             
-            # 获取响应内容
-            response_content = response.text
-            print(f"Response status: {response.status_code}")
-            print(f"Response content: {response_content}")
-            
-            # 尝试解析JSON响应
-            try:
-                return Response(
-                    content=response_content,
-                    status_code=response.status_code,
-                    headers=dict(response.headers)
-                )
-            except json.JSONDecodeError as e:
-                print(f"JSON decode error: {str(e)}")
-                return Response(
-                    content=response_content,
-                    status_code=response.status_code,
-                    headers=dict(response.headers)
-                )
+            return Response(
+                content=response.text,
+                status_code=response.status_code,
+                headers=dict(response.headers)
+            )
                 
         except Exception as e:
             print(f"Proxy Error: {str(e)}")