Update app.py

app.py

@@ -13,6 +13,8 @@ import threading
 import aiohttp
 import asyncio
 from quart import Quart, request, jsonify
+from flask_socketio import SocketIO, emit
+import eventlet
 
 # Create a queue for log messages
 log_queue = queue.Queue()
@@ -30,16 +32,23 @@ queue_handler.setFormatter(logging.Formatter('%(asctime)s - %(levelname)s - %(me
 logger.addHandler(queue_handler)
 logger.setLevel(logging.INFO)
 
+# Initialize Flask app with eventlet
+eventlet.monkey_patch()
 app = Flask(__name__)
-# Enable CORS with specific settings
+
+# Configure CORS with WebSocket support
 CORS(app, resources={
     r"/*": {
-        "origins": "*",
+        "origins": ["http://localhost:8000", "http://127.0.0.1:8000"],  # Add your Laravel domains
         "methods": ["GET", "POST", "OPTIONS"],
-        "allow_headers": ["Content-Type", "Authorization"]
+        "allow_headers": ["Content-Type", "Authorization"],
+        "supports_credentials": True
     }
 })
 
+# Initialize SocketIO with CORS support
+socketio = SocketIO(app, cors_allowed_origins=["http://localhost:8000", "http://127.0.0.1:8000"], async_mode='eventlet')
+
 # Configuration
 OPENAI_API_KEY = os.getenv('OPENAI_API_KEY')
 OPENROUTER_API_KEY = os.getenv('OPENROUTER_API_KEY')
@@ -48,6 +57,39 @@ SERPAPI_API_KEY = os.getenv('SERPAPI_API_KEY')
 
 ssl._create_default_https_context = ssl._create_unverified_context
 
+# WebSocket event handlers
+@socketio.on('connect')
+def handle_connect():
+    logger.info("Client connected")
+    emit('connection_response', {'status': 'connected'})
+
+@socketio.on('disconnect')
+def handle_disconnect():
+    logger.info("Client disconnected")
+
+@socketio.on_error()
+def error_handler(e):
+    logger.error(f"SocketIO error: {str(e)}")
+    emit('error', {'error': str(e)})
+
+# Modified stream_logs route to use WebSocket
+@socketio.on('subscribe_logs')
+def handle_subscribe_logs():
+    try:
+        while True:
+            try:
+                log_message = log_queue.get(timeout=1)
+                emit('log_message', {'message': log_message})
+            except queue.Empty:
+                emit('heartbeat', {'message': 'heartbeat'})
+            except Exception as e:
+                logger.error(f"Error in log streaming: {e}")
+                emit('error', {'error': str(e)})
+            eventlet.sleep(0.1)  # Prevent CPU overuse
+    except Exception as e:
+        logger.error(f"Error in log subscription: {e}")
+        emit('error', {'error': str(e)})
+
 def generate_preliminary_plan(cluster_data):
     logger.info("Generating preliminary plan...")
     try:
@@ -601,32 +643,6 @@ async def generate_from_csv_text():
         logger.error(f"Error in generate_from_csv_text: {e}")
         return jsonify({'error': str(e)}), 500
 
-@app.route('/logs/stream')
-def stream_logs():
-    def generate():
-        while True:
-            try:
-                # Get log message from queue, timeout after 1 second
-                log_message = log_queue.get(timeout=1)
-                yield f"data: {log_message}\n\n"
-            except queue.Empty:
-                # Send a heartbeat to keep the connection alive
-                yield "data: heartbeat\n\n"
-            except GeneratorExit:
-                break
-
-    response = Response(stream_with_context(generate()), mimetype='text/event-stream')
-    response.headers['Cache-Control'] = 'no-cache'
-    response.headers['Connection'] = 'keep-alive'
-    response.headers['Access-Control-Allow-Origin'] = '*'
-    return response
-
 if __name__ == '__main__':
-    logger.info("Starting Flask API server...")
-    import hypercorn.asyncio
-    import hypercorn.config
-
-    config = hypercorn.config.Config()
-    config.bind = ["127.0.0.1:5001"]
-
-    asyncio.run(hypercorn.asyncio.serve(app, config))
+    logger.info("Starting Flask API server with SocketIO...")
+    socketio.run(app, host='127.0.0.1', port=5001, debug=True)