Update app.py

app.py

@@ -1093,104 +1093,93 @@ async def websocket_chat(ws: WebSocket):
         # Internal endpoint (avoids Hugging Face proxy redirect)
         internal_url = "http://127.0.0.1:7860/gen/chat/completions"
 
-        # Queue to ensure sequential request processing
-        import asyncio
-        request_queue = asyncio.Queue()
-        
-        async def process_requests():
-            # Persistent HTTP client for streaming
-            async with httpx.AsyncClient(
-                timeout=None,
-                follow_redirects=False
-            ) as client:
+        # Persistent HTTP client for streaming
+        async with httpx.AsyncClient(
+            timeout=None,
+            follow_redirects=False
+        ) as client:
+            request_counter = 0
+            
+            # Background task to handle incoming requests
+            async def handle_incoming_requests():
+                nonlocal request_counter
                 while True:
                     try:
-                        request_id, body, headers = await request_queue.get()
-                    except asyncio.CancelledError:
+                        msg = await ws.receive_text()
+                    except WebSocketDisconnect:
                         break
 
                     try:
-                        async with client.stream(
-                            "POST",
-                            internal_url,
-                            json=body,
-                            headers=headers
-                        ) as response:
-
-                            # Handle upstream errors
-                            if response.status_code >= 400:
-                                error_text = ""
-                                try:
-                                    error_text = (await response.aread()).decode(
-                                        "utf-8", errors="replace"
-                                    )[:500]
-                                except Exception:
-                                    pass
-
-                                error_payload = json.dumps({
-                                    "error": "Upstream request failed",
-                                    "status": response.status_code,
-                                    "detail": error_text
-                                })
-                                try:
-                                    await ws.send_text(f"{request_id}:{error_payload}")
-                                except RuntimeError:
-                                    return
-                                continue
-
-                            # Stream tokens/lines to the websocket
-                            async for line in response.aiter_lines():
-                                if not line:
-                                    continue
+                        data = json.loads(msg)
+                    except Exception:
+                        await ws.send_json({"error": "Invalid JSON"})
+                        continue
 
-                                try:
-                                    await ws.send_text(f"{request_id}:{line}")
-                                except RuntimeError:
-                                    return
+                    body = data.get("body")
+                    headers = data.get("headers") or {}
 
-                    except Exception as stream_error:
-                        try:
-                            error_payload = json.dumps({
-                                "error": "Streaming error",
-                                "detail": str(stream_error)
-                            })
-                            await ws.send_text(f"{request_id}:{error_payload}")
-                        except Exception:
-                            pass
+                    if not body:
+                        await ws.send_json({"error": "Missing body"})
+                        continue
 
-        # Start the request processor task
-        processor_task = asyncio.create_task(process_requests())
+                    request_counter += 1
+                    request_id = request_counter
+                    
+                    # Process this request concurrently
+                    asyncio.create_task(process_request(client, request_id, body, headers))
 
-        try:
-            request_counter = 0
-            while True:
+            async def process_request(client, request_id, body, headers):
                 try:
-                    msg = await ws.receive_text()
-                except WebSocketDisconnect:
-                    break
+                    async with client.stream(
+                        "POST",
+                        internal_url,
+                        json=body,
+                        headers=headers
+                    ) as response:
 
-                try:
-                    data = json.loads(msg)
-                except Exception:
-                    await ws.send_json({"error": "Invalid JSON"})
-                    continue
+                        # Handle upstream errors
+                        if response.status_code >= 400:
+                            error_text = ""
+                            try:
+                                error_text = (await response.aread()).decode(
+                                    "utf-8", errors="replace"
+                                )[:500]
+                            except Exception:
+                                pass
 
-                body = data.get("body")
-                headers = data.get("headers") or {}
+                            error_payload = json.dumps({
+                                "error": "Upstream request failed",
+                                "status": response.status_code,
+                                "detail": error_text
+                            })
+                            try:
+                                await ws.send_text(f"{request_id}:{error_payload}")
+                            except RuntimeError:
+                                return
+                            return
 
-                if not body:
-                    await ws.send_json({"error": "Missing body"})
-                    continue
+                        # Stream tokens/lines to the websocket
+                        async for line in response.aiter_lines():
+                            if not line:
+                                continue
 
-                request_counter += 1
-                await request_queue.put((request_counter, body, headers))
+                            try:
+                                await ws.send_text(f"{request_id}:{line}")
+                            except RuntimeError:
+                                return
 
-        finally:
-            processor_task.cancel()
-            try:
-                await processor_task
-            except asyncio.CancelledError:
-                pass
+                except Exception as stream_error:
+                    try:
+                        error_payload = json.dumps({
+                            "error": "Streaming error",
+                            "detail": str(stream_error)
+                        })
+                        await ws.send_text(f"{request_id}:{error_payload}")
+                    except Exception:
+                        pass
+
+            # Start handling incoming requests
+            await handle_incoming_requests()
 
     except WebSocketDisconnect:
         return