Update main.py

main.py

@@ -5,7 +5,7 @@ import socket
 
 app = FastAPI()
 
-# Replace this with your actual backend API key
+# Your real backend API key
 REAL_API_KEY = "sk-kwWVUQPDsLemvilLcyzSSWRzo8sctCzxzlbdN0ZC5ZUCCv0m"
 BASE_URL = "https://fast.typegpt.net"
 PUBLIC_AUTH_TOKEN = "TypeGPT-Free4ALL"
@@ -22,16 +22,14 @@ async def startup_event():
 
 @app.api_route("/{path:path}", methods=["GET", "POST"])
 async def proxy(request: Request, path: str, authorization: str = Header(None)):
-    # Validate token
     if not authorization or not authorization.startswith("Bearer "):
         raise HTTPException(status_code=401, detail="Missing or malformed Authorization header.")
-    
+
     token = authorization.replace("Bearer ", "").strip()
     if token != PUBLIC_AUTH_TOKEN:
         raise HTTPException(status_code=401, detail="Invalid Authorization token. Use 'TypeGPT-Free4ALL'.")
 
     target_url = f"{BASE_URL}/{path}"
-
     headers = {
         "Authorization": f"Bearer {REAL_API_KEY}",
         "Content-Type": request.headers.get("content-type", "application/json"),
@@ -46,42 +44,46 @@ async def proxy(request: Request, path: str, authorization: str = Header(None)):
     print(f"📦 Request Body (first 200 chars): {body[:200]}")
     print(f"🔁 Stream mode: {is_stream}")
 
-    async with httpx.AsyncClient(timeout=60) as client:
-        try:
-            if is_stream:
-                async def stream_generator():
-                    async with client.stream(
-                        method=request.method,
-                        url=target_url,
-                        headers=headers,
-                        content=body
-                    ) as upstream_response:
-                        async for chunk in upstream_response.aiter_bytes():
-                            yield chunk
+    if is_stream:
+        # Move the client and stream both inside the generator
+        async def stream_generator():
+            async with httpx.AsyncClient(timeout=60) as client:
+                async with client.stream(
+                    method=request.method,
+                    url=target_url,
+                    headers=headers,
+                    content=body
+                ) as upstream_response:
+                    async for chunk in upstream_response.aiter_bytes():
+                        yield chunk
 
-                return StreamingResponse(
-                    stream_generator(),
-                    status_code=200,
-                    media_type="text/event-stream"
-                )
-            else:
+        return StreamingResponse(
+            stream_generator(),
+            status_code=200,
+            media_type="text/event-stream"
+        )
+
+    else:
+        # Non-stream (normal JSON response)
+        async with httpx.AsyncClient(timeout=60) as client:
+            try:
                 response = await client.request(
                     method=request.method,
                     url=target_url,
                     headers=headers,
                     content=body
                 )
-        except httpx.RequestError as e:
-            raise HTTPException(status_code=502, detail=f"Backend request failed: {e}")
+            except httpx.RequestError as 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]}")
+        print(f"↩️ TypeGPT Status: {response.status_code}")
+        print(f"🧾 Response Snippet: {response.text[:200]}")
 
-    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")
-        )
+        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")
+            )