feat: add websocket server for real-time communication

app.py

@@ -1,12 +1,16 @@
+import asyncio
+import websockets
+import json
+import threading
 from flask import Flask, request, jsonify
 import numpy as np
-import json
-from .music_generator import MusicGenerator
+from music_generator import MusicGenerator
 
+# --- Existing Flask App Setup ---
 app = Flask(__name__)
 
 # Load the consonance matrix
-with open('/home/KidIkaros/Documents/code/Ikaros/musick/chord_detector_extension/consonance_matrix.json') as f:
+with open('consonance_matrix.json') as f:
     consonance_matrix = np.array(json.load(f))
 
 notes = ['C', 'C#', 'D', 'D#', 'E', 'F', 'F#', 'G', 'G#', 'A', 'A#', 'B']
@@ -18,111 +22,127 @@ def note_to_index(note):
 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})
 
-@app.route('/generate', methods=['POST'])
-def generate():
-    data = request.get_json()
-    start_sequence_indices = [note_to_index(note) for note in data.get('start_sequence', [])]
-    length = data.get('length', 10)
-
-    if not start_sequence_indices:
-        return jsonify({'generated_sequence': []})
-
-    generated_indices = generator.generate(start_sequence_indices, length)
-    generated_notes = [index_to_note(i) for i in generated_indices]
-    return jsonify({'generated_sequence': generated_notes})
-
-
-@app.route('/analyze_harmony', methods=['POST'])
-def analyze_harmony():
-    data = request.get_json()
-    history = data.get('history', [])
-
-    if len(history) < 2:
-        return jsonify({'harmony_scores': []})
-
-    harmony_scores = []
-    for i in range(len(history) - 1):
-        try:
-            note1_index = note_to_index(history[i]['chord'])
-            note2_index = note_to_index(history[i+1]['chord'])
-            score = consonance_matrix[note1_index, note2_index]
-            harmony_scores.append(score)
-        except (ValueError, IndexError):
-            # Handle cases where a chord is not in our 'notes' list (e.g., 'N')
-            harmony_scores.append(0) # Assign a neutral score
 
-    return jsonify({'harmony_scores': harmony_scores})
-
-@app.route('/what_if', methods=['POST'])
-def what_if():
-    data = request.get_json()
-    history = data.get('history', [])
-    suggestion_index = data.get('suggestion')
-
-    if len(history) < 1 or suggestion_index is None:
-        return jsonify({'harmony_score': 0})
+# --- 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"]])
+
+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)
+    key = "C Major" # Mock key
+
+    analysis_result = {
+        "type": "analysis_update",
+        "current_chord": detected_chord,
+        "predicted_chord": predicted_chord,
+        "musical_key": key
+    }
+    print(f"Broadcasting analysis: {analysis_result}")
+    await broadcast_to_webapps(json.dumps(analysis_result))
+
+
+async def connection_handler(websocket, path):
+    """Handles incoming WebSocket connections."""
+    global audio_source
+    print(f"New client connected.")
 
     try:
-        last_note_index = note_to_index(history[-1]['chord'])
-        score = consonance_matrix[last_note_index, suggestion_index]
-    except (ValueError, IndexError):
-        score = 0
-
-    return jsonify({'harmony_score': score})
-
-
-from sklearn.decomposition import PCA
-
-@app.route('/song_fingerprint', methods=['POST'])
-def song_fingerprint():
-    data = request.get_json()
-    history = data.get('history', [])
-
-    if len(history) < 3:
-        return jsonify({'fingerprint': []})
-
-    # Create a matrix of chord transitions
-    transitions = []
-    for i in range(len(history) - 1):
-        try:
-            note1_index = note_to_index(history[i]['chord'])
-            note2_index = note_to_index(history[i+1]['chord'])
-            transitions.append([note1_index, note2_index])
-        except (ValueError, IndexError):
-            pass
-
-    if len(transitions) < 3:
-        return jsonify({'fingerprint': []})
-
-    # Use PCA to reduce to 3 dimensions
-    pca = PCA(n_components=3)
-    fingerprint = pca.fit_transform(transitions).tolist()
-
-    return jsonify({'fingerprint': fingerprint})
-
-
-if __name__ == '__main__':
-    # Train the generator on some dummy data
-    sequences = [
-        [0, 4, 7, 0], # Cmaj -> C
-        [5, 9, 0, 5]  # Fmaj -> F
-    ]
-    generator.train(sequences)
-    app.run(port=5000)
+        # The first message from a client identifies its role.
+        initial_message = await websocket.recv()
+        message_data = json.loads(initial_message)
+        client_type = message_data.get("type")
+
+        if client_type == "extension_hello":
+            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"}))
+
+        elif client_type == "webapp_hello":
+            clients["webapp"].add(websocket)
+            print("Web app client connected.")
+            await websocket.send(json.dumps({"status": "connected", "role": "viewer"}))
+
+        else:
+            print(f"Unknown client type: {client_type}. Disconnecting.")
+            return
+
+        # Listen for messages from the client
+        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.")
+    finally:
+        # Remove the client from our sets upon disconnection
+        if websocket in clients["webapp"]:
+            clients["webapp"].remove(websocket)
+        if websocket == audio_source:
+            audio_source = None
+            print("Audio capture extension 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.
+    websocket_port = 7860
+    print(f"Starting WebSocket server on port {websocket_port}...")
+    start_server = websockets.serve(connection_handler, "0.0.0.0", websocket_port)
+
+    asyncio.get_event_loop().run_until_complete(start_server)
+    asyncio.get_event_loop().run_forever()

requirements.txt

@@ -2,4 +2,5 @@ networkx==3.3
 numpy==1.26.4
 flask
 flask-cors
-tensorflow
+tensorflow
+websockets