refactor: simplify to pure websocket server

app.py

@@ -1,13 +1,10 @@
 import asyncio
 import websockets
 import json
-import threading
-from flask import Flask, request, jsonify
 import numpy as np
 from music_generator import MusicGenerator
 
-# --- Existing Flask App Setup ---
-app = Flask(__name__)
+# --- WebSocket Server Setup ---
 
 # Load the consonance matrix
 with open('consonance_matrix.json') as f:
@@ -16,53 +13,32 @@ with open('consonance_matrix.json') as f:
 notes = ['C', 'C#', 'D', 'D#', 'E', 'F', 'F#', 'G', 'G#', 'A', 'A#', 'B']
 generator = MusicGenerator(len(notes))
 
-def note_to_index(note):
-    return notes.index(note.split(':')[0])
-
-def index_to_note(index):
-    return notes[index]
-
-# (Keep all the existing @app.route endpoints for now)
-@app.route('/predict', methods=['POST'])
-def predict():
-    # This route will likely be deprecated in favor of WebSockets
-    # but we keep it for now.
-    data = request.get_json()
-    history = data.get('history', [])
-    if len(history) < 1:
-        return jsonify({'prediction': 'N/A'})
-    try:
-        last_note_index = note_to_index(history[-1]['chord'])
-        prediction_index = generator.generate([last_note_index], length=1)[-1]
-        prediction = index_to_note(prediction_index)
-    except (ValueError, IndexError):
-        prediction = 'N/A'
-    return jsonify({'prediction': prediction})
-
-
-# --- WebSocket Server Setup ---
-
 # In-memory storage for connected clients
-# We'll have two types of clients: 'extension' and 'webapp'
 clients = {
     "webapp": set()
 }
-# We only need one audio source, so we don't need a set for the extension.
 audio_source = None
 
 async def broadcast_to_webapps(message):
     """Sends a message to all connected webapp clients."""
     if clients["webapp"]:
-        await asyncio.wait([client.send(message) for client in clients["webapp"]])
+        # Create a copy of the set to avoid issues if a client disconnects during broadcast
+        disconnected_clients = set()
+        for client in clients["webapp"]:
+            try:
+                await client.send(message)
+            except websockets.exceptions.ConnectionClosed:
+                disconnected_clients.add(client)
+        # Remove clients that have disconnected
+        for client in disconnected_clients:
+            clients["webapp"].remove(client)
+
 
 async def handle_audio_data(data):
     """
     This is the core audio processing function.
     For now, it will just mock the analysis.
-    In the future, this is where we'll plug in our TensorFlow model.
     """
-    # Mock analysis: Pretend we detected a chord and generated a prediction.
-    # We can make this more interesting by picking a random chord.
     import random
     detected_chord = random.choice(notes)
     predicted_chord = random.choice(notes)
@@ -90,15 +66,26 @@ async def connection_handler(websocket, path):
         client_type = message_data.get("type")
 
         if client_type == "extension_hello":
+            if audio_source is not None:
+                # If there's already an extension connected, disconnect the old one.
+                print("An extension is already connected. Disconnecting the old one.")
+                await audio_source.close(reason="New extension connected.")
+            
             audio_source = websocket
-            clients["webapp"].add(websocket) # Also treat extension as a webapp to receive messages
             print("Audio capture extension connected.")
             await websocket.send(json.dumps({"status": "connected", "role": "audio_source"}))
+            await broadcast_to_webapps(json.dumps({"type": "status_update", "message": "Audio source connected."}))
+
 
         elif client_type == "webapp_hello":
             clients["webapp"].add(websocket)
             print("Web app client connected.")
             await websocket.send(json.dumps({"status": "connected", "role": "viewer"}))
+            if audio_source is not None:
+                 await websocket.send(json.dumps({"type": "status_update", "message": "Audio source connected."}))
+            else:
+                 await websocket.send(json.dumps({"type": "status_update", "message": "Waiting for audio source..."}))
+
 
         else:
             print(f"Unknown client type: {client_type}. Disconnecting.")
@@ -108,12 +95,10 @@ async def connection_handler(websocket, path):
         async for message in websocket:
             if websocket == audio_source:
                 # This is audio data from the extension
-                # For now, we assume the message is a chunk of audio data.
-                # We will simply trigger our mock analysis.
                 await handle_audio_data(message)
 
     except websockets.exceptions.ConnectionClosed:
-        print("Client disconnected.")
+        print(f"Client disconnected: {websocket.remote_address}")
     finally:
         # Remove the client from our sets upon disconnection
         if websocket in clients["webapp"]:
@@ -121,28 +106,15 @@ async def connection_handler(websocket, path):
         if websocket == audio_source:
             audio_source = None
             print("Audio capture extension disconnected.")
+            await broadcast_to_webapps(json.dumps({"type": "status_update", "message": "Audio source disconnected."}))
 
 
-def run_flask_app():
-    """Runs the Flask app in a separate thread."""
-    # Note: Using Flask's development server is not ideal for production.
-    # A proper WSGI server like Gunicorn should be used.
-    # But for Hugging Face Spaces, this is often sufficient.
-    app.run(host='0.0.0.0', port=5000)
-
-
-if __name__ == "__main__":
-    # Start the Flask app in a background thread
-    flask_thread = threading.Thread(target=run_flask_app)
-    flask_thread.daemon = True
-    flask_thread.start()
-
-    # Start the WebSocket server
-    # Hugging Face Spaces exposes port 7860 by default for web traffic.
-    # We will use this port for our WebSocket server.
+async def main():
+    """Starts the WebSocket server."""
     websocket_port = 7860
-    print(f"Starting WebSocket server on port {websocket_port}...")
-    start_server = websockets.serve(connection_handler, "0.0.0.0", websocket_port)
+    print(f"Starting pure WebSocket server on port {websocket_port}...")
+    async with websockets.serve(connection_handler, "0.0.0.0", websocket_port):
+        await asyncio.Future()  # run forever
 
-    asyncio.get_event_loop().run_until_complete(start_server)
-    asyncio.get_event_loop().run_forever()
+if __name__ == "__main__":
+    asyncio.run(main())

requirements.txt

@@ -1,6 +1,4 @@
 networkx==3.3
 numpy==1.26.4
-flask
-flask-cors
 tensorflow
 websockets