Update main.py

main.py

@@ -1,18 +1,15 @@
 from fastapi import FastAPI, Request, Header, HTTPException
-from fastapi.responses import JSONResponse, Response
+from fastapi.responses import JSONResponse, Response, StreamingResponse
 import httpx
 import socket
+import asyncio
 
 app = FastAPI()
 
-# Real secret API key (used internally, never exposed)
-REAL_API_KEY = "sk-94NDKhKQkhKoYnY65mg4NAIFK5BqNiCtxo8u3PsDpb0IucZt"
+REAL_API_KEY = "sk-94NDKhKQkhKoYnY65mg4NAIFK5BqNiCtxo8u3PsDpb0IucZt"  # Replace this with your actual key
 BASE_URL = "https://fast.typegpt.net"
-
-# Public key that users must use in Authorization header
 PUBLIC_AUTH_TOKEN = "TypeGPT-Free4ALL"
 
-# Show server IP at startup
 @app.on_event("startup")
 async def startup_event():
     try:
@@ -25,7 +22,6 @@ async def startup_event():
 
 @app.api_route("/{path:path}", methods=["GET", "POST"])
 async def proxy(request: Request, path: str, authorization: str = Header(None)):
-    # Validate Authorization token
     if not authorization or not authorization.startswith("Bearer "):
         raise HTTPException(status_code=401, detail="Missing or malformed Authorization header.")
 
@@ -33,10 +29,8 @@ async def proxy(request: Request, path: str, authorization: str = Header(None)):
     if token != PUBLIC_AUTH_TOKEN:
         raise HTTPException(status_code=401, detail="Invalid Authorization token. Use 'TypeGPT-Free4ALL'.")
 
-    # Rebuild the target URL
     target_url = f"{BASE_URL}/{path}"
 
-    # Use clean headers to avoid Cloudflare issues
     headers = {
         "Authorization": f"Bearer {REAL_API_KEY}",
         "Content-Type": request.headers.get("content-type", "application/json"),
@@ -44,34 +38,40 @@ async def proxy(request: Request, path: str, authorization: str = Header(None)):
         "User-Agent": "FastAPI-Proxy"
     }
 
-    # Read the body
     body = await request.body()
 
-    # Log for debugging
+    # Detect stream mode
+    is_stream = b'"stream":true' in body or b'"stream": true' in body
+
     print(f"🔄 Forwarding to: {target_url}")
     print(f"📦 Request Body (first 200 chars): {body[:200]}")
+    print(f"🔁 Stream mode: {is_stream}")
 
-    async with httpx.AsyncClient() as client:
+    async with httpx.AsyncClient(timeout=60) as client:
         try:
             response = await client.request(
                 method=request.method,
                 url=target_url,
                 headers=headers,
                 content=body,
-                timeout=60
+                stream=is_stream,
             )
         except httpx.RequestError as e:
-            raise HTTPException(status_code=502, detail=f"Failed to reach backend: {e}")
+            raise HTTPException(status_code=502, detail=f"Backend request failed: {e}")
 
-    print(f"↩️ TypeGPT Status: {response.status_code}")
-    print(f"🧾 Response Snippet: {response.text[:200]}")
+        if is_stream:
+            async def event_generator():
+                async for chunk in response.aiter_bytes():
+                    yield chunk
 
-    # Return JSON if possible, else fallback
-    try:
-        return JSONResponse(content=response.json(), status_code=response.status_code)
-    except Exception:
-        return Response(
-            content=response.content,
-            status_code=response.status_code,
-            media_type=response.headers.get("content-type", "text/plain")
-        )
+            return StreamingResponse(event_generator(), status_code=response.status_code, media_type="text/event-stream")
+
+        # For non-streamed response
+        try:
+            return JSONResponse(content=response.json(), status_code=response.status_code)
+        except Exception:
+            return Response(
+                content=response.content,
+                status_code=response.status_code,
+                media_type=response.headers.get("content-type", "text/plain")
+            )