alchemist-nrel 0.3.0__py3-none-any.whl → 0.3.2__py3-none-any.whl

api/routers/visualizations.py

@@ -194,9 +194,10 @@ async def get_contour_data(
     # CRITICAL FIX: Reorder columns to match training data
     # The model was trained with a specific column order, we must match it.
     # Exclude metadata columns that are part of the experiments table but
-    # are not model input features (e.g., Iteration, Reason, Output, Noise).
+    # are not model input features (e.g., Iteration, Reason, target columns, Noise).
     train_data = session.experiment_manager.get_data()
-    metadata_cols = {'Iteration', 'Reason', 'Output', 'Noise'}
+    target_cols = set(session.experiment_manager.target_columns)
+    metadata_cols = {'Iteration', 'Reason', 'Noise'} | target_cols
     feature_cols = [col for col in train_data.columns if col not in metadata_cols]
 
     # Safely align the prediction grid to the model feature order.
@@ -216,11 +217,12 @@ async def get_contour_data(
     experiments_data = None
     if request.include_experiments and len(session.experiment_manager) > 0:
         exp_df = session.experiment_manager.get_data()
-        if request.x_var in exp_df.columns and request.y_var in exp_df.columns and "Output" in exp_df.columns:
+        target_col = session.experiment_manager.target_columns[0]  # Use first target for visualization
+        if request.x_var in exp_df.columns and request.y_var in exp_df.columns and target_col in exp_df.columns:
             experiments_data = {
                 "x": exp_df[request.x_var].tolist(),
                 "y": exp_df[request.y_var].tolist(),
-                "output": exp_df["Output"].tolist()
+                "output": exp_df[target_col].tolist()
             }
     
     # Get suggestion data if requested (would need to be stored in session)

api/routers/websocket.py

@@ -0,0 +1,132 @@
+"""
+WebSocket router for real-time session updates.
+
+Provides real-time push notifications for session events like lock status changes,
+eliminating the need for client-side polling.
+"""
+
+from fastapi import APIRouter, WebSocket, WebSocketDisconnect
+from typing import Dict, Set
+import json
+import logging
+
+logger = logging.getLogger(__name__)
+
+router = APIRouter()
+
+# Track active WebSocket connections per session
+# Structure: {session_id: {websocket1, websocket2, ...}}
+active_connections: Dict[str, Set[WebSocket]] = {}
+
+
+@router.websocket("/ws/sessions/{session_id}")
+async def websocket_endpoint(websocket: WebSocket, session_id: str):
+    """
+    WebSocket endpoint for real-time session updates.
+    
+    Clients connect to this endpoint to receive push notifications about
+    session events (lock status changes, experiment additions, etc.).
+    
+    Args:
+        websocket: WebSocket connection
+        session_id: Session ID to subscribe to
+    """
+    await websocket.accept()
+    logger.debug(f"WebSocket connected: session_id={session_id}")
+    
+    # Register this connection for this session
+    if session_id not in active_connections:
+        active_connections[session_id] = set()
+    active_connections[session_id].add(websocket)
+    
+    try:
+        # Send initial connection confirmation
+        await websocket.send_json({
+            "event": "connected",
+            "session_id": session_id,
+            "message": "WebSocket connection established"
+        })
+        
+        # Keep connection alive and listen for client messages
+        while True:
+            # Receive messages from client (for future bi-directional features)
+            data = await websocket.receive_text()
+            
+            # Echo back for debugging (can be removed in production)
+            try:
+                message = json.loads(data)
+                logger.debug(f"Received from client: {message}")
+            except json.JSONDecodeError:
+                logger.warning(f"Invalid JSON from client: {data}")
+            
+    except WebSocketDisconnect:
+        logger.debug(f"WebSocket disconnected: session_id={session_id}")
+    finally:
+        # Clean up on disconnect
+        if session_id in active_connections:
+            active_connections[session_id].discard(websocket)
+            
+            # Remove session entry if no more connections
+            if not active_connections[session_id]:
+                del active_connections[session_id]
+                logger.debug(f"No more connections for session {session_id}")
+
+
+async def broadcast_to_session(session_id: str, event: dict):
+    """
+    Broadcast an event to all WebSocket clients connected to a session.
+    
+    Args:
+        session_id: Session ID to broadcast to
+        event: Event data to send (will be JSON serialized)
+    """
+    if session_id not in active_connections:
+        logger.debug(f"No active connections for session {session_id}")
+        return
+    
+    # Track dead connections
+    dead_connections = set()
+    
+    # Send to all connected clients
+    for connection in active_connections[session_id]:
+        try:
+            await connection.send_json(event)
+            logger.debug(f"Broadcast to session {session_id}: {event.get('event')}")
+        except Exception as e:
+            logger.warning(f"Failed to send to connection: {e}")
+            dead_connections.add(connection)
+    
+    # Clean up dead connections
+    if dead_connections:
+        active_connections[session_id] -= dead_connections
+        logger.info(f"Cleaned up {len(dead_connections)} dead connections")
+        
+        # Remove session if no connections left
+        if not active_connections[session_id]:
+            del active_connections[session_id]
+
+
+def get_connection_count(session_id: str) -> int:
+    """
+    Get the number of active WebSocket connections for a session.
+    
+    Args:
+        session_id: Session ID to check
+        
+    Returns:
+        Number of active connections
+    """
+    return len(active_connections.get(session_id, set()))
+
+
+def get_all_connection_counts() -> Dict[str, int]:
+    """
+    Get connection counts for all sessions.
+    
+    Returns:
+        Dictionary mapping session_id to connection count
+    """
+    return {
+        session_id: len(connections)
+        for session_id, connections in active_connections.items()
+    }

run_api.py → api/run_api.py

@@ -2,10 +2,10 @@
 Startup script for ALchemist FastAPI server.
 
 Usage:
-    python run_api.py              # Development mode with auto-reload
-    python run_api.py --production # Production mode (no reload)
-    python run_api.py --dev        # Explicitly start in development mode
-    alchemist-web                  # Entry point (production mode by default)
+    python -m api.run_api              # Development mode with auto-reload
+    python -m api.run_api --production # Production mode (no reload)
+    python -m api.run_api --dev        # Explicitly start in development mode
+    alchemist-web                      # Entry point (production mode by default)
 """
 
 def main():
@@ -15,10 +15,10 @@ def main():
     
     # For the alchemist-web entry point, default to production mode
     # Only use dev mode if explicitly requested
-    is_script_call = any(arg.endswith('run_api.py') for arg in sys.argv)
+    is_script_call = any(arg.endswith('run_api.py') or 'api.run_api' in arg for arg in sys.argv)
     
     if is_script_call:
-        # Called as: python run_api.py
+        # Called as: python -m api.run_api
         # Default to dev mode unless --production flag is present
         production = "--production" in sys.argv or "--prod" in sys.argv
     else:
@@ -47,7 +47,8 @@ def main():
             host="0.0.0.0",
             port=8000,
             reload=True,
-            log_level="info"
+            log_level="warning",  # Suppress INFO logs from polling
+            access_log=False  # Disable access logs entirely
         )