alchemist-nrel 0.2.1__py3-none-any.whl → 0.3.1__py3-none-any.whl

api/routers/sessions.py

@@ -2,13 +2,24 @@
 Sessions router - Session lifecycle management.
 """
 
-from fastapi import APIRouter, HTTPException, status, UploadFile, File
-from fastapi.responses import Response
-from ..models.responses import SessionCreateResponse, SessionInfoResponse
+from fastapi import APIRouter, HTTPException, status, UploadFile, File, Depends
+from fastapi.responses import Response, FileResponse, JSONResponse
+from typing import Optional
+from ..models.requests import UpdateMetadataRequest, LockDecisionRequest, SessionLockRequest
+from ..models.responses import (
+    SessionCreateResponse, SessionInfoResponse, SessionStateResponse,
+    SessionMetadataResponse, AuditLogResponse, AuditEntryResponse, LockDecisionResponse,
+    SessionLockResponse
+)
+from .websocket import broadcast_to_session
 from ..services import session_store
 from ..dependencies import get_session
+from alchemist_core.session import OptimizationSession
 from datetime import datetime
 import logging
+import json
+from pathlib import Path
+import tempfile
 
 logger = logging.getLogger(__name__)
 
@@ -50,6 +61,36 @@ async def get_session_info(session_id: str):
     return SessionInfoResponse(**info)
 
 
+@router.get("/sessions/{session_id}/state", response_model=SessionStateResponse)
+async def get_session_state(
+    session_id: str,
+    session: OptimizationSession = Depends(get_session)
+):
+    """
+    Get current session state for monitoring autonomous optimization.
+    
+    Returns key metrics for dashboard displays or autonomous controllers
+    to monitor optimization progress without retrieving full session data.
+    """
+    # Get session metrics
+    n_variables = len(session.search_space.variables)
+    n_experiments = len(session.experiment_manager.df)
+    model_trained = session.model is not None
+    
+    # Get last suggestion if available
+    last_suggestion = None
+    if hasattr(session, '_last_suggestion') and session._last_suggestion:
+        last_suggestion = session._last_suggestion
+    
+    return SessionStateResponse(
+        session_id=session_id,
+        n_variables=n_variables,
+        n_experiments=n_experiments,
+        model_trained=model_trained,
+        last_suggestion=last_suggestion
+    )
+
+
 @router.delete("/sessions/{session_id}", status_code=status.HTTP_204_NO_CONTENT)
 async def delete_session(session_id: str):
     """
@@ -89,13 +130,30 @@ async def extend_session(session_id: str, hours: int = 24):
     }
 
 
+@router.post("/sessions/{session_id}/save", status_code=status.HTTP_200_OK)
+async def save_session_server_side(session_id: str):
+    """
+    Persist the current in-memory session to the server-side session file.
+
+    This allows the web UI to save changes directly to the session store file
+    instead of triggering a browser download.
+    """
+    success = session_store.persist_session_to_disk(session_id)
+    if not success:
+        raise HTTPException(
+            status_code=status.HTTP_404_NOT_FOUND,
+            detail=f"Session {session_id} not found or failed to save"
+        )
+    return {"message": "Session persisted to server storage"}
+
+
 @router.get("/sessions/{session_id}/export")
 async def export_session(session_id: str):
     """
     Export a session for download.
     
-    Downloads the complete session state as a .pkl file that can be
-    reimported later.
+    Downloads the complete session state as a .json file that can be
+    reimported later or used in desktop application.
     """
     session_data = session_store.export_session(session_id)
     if session_data is None:
@@ -106,9 +164,9 @@ async def export_session(session_id: str):
     
     return Response(
         content=session_data,
-        media_type="application/octet-stream",
+        media_type="application/json",
         headers={
-            "Content-Disposition": f"attachment; filename=session_{session_id}.pkl"
+            "Content-Disposition": f"attachment; filename=session_{session_id}.json"
         }
     )
 
@@ -118,12 +176,14 @@ async def import_session(file: UploadFile = File(...)):
     """
     Import a previously exported session.
     
-    Uploads a .pkl session file and creates a new session with the imported data.
-    A new session ID will be generated.
+    Uploads a .json session file and creates a new session with the imported data.
+    A new session ID will be generated. Compatible with desktop application sessions.
     """
     try:
         session_data = await file.read()
-        session_id = session_store.import_session(session_data)
+        # Decode bytes to string for JSON
+        session_json = session_data.decode('utf-8')
+        session_id = session_store.import_session(session_json)
         
         if session_id is None:
             raise HTTPException(
@@ -144,3 +204,371 @@ async def import_session(file: UploadFile = File(...)):
             status_code=status.HTTP_400_BAD_REQUEST,
             detail=f"Failed to import session: {str(e)}"
         )
+
+
+# ============================================================
+# Metadata Management Endpoints
+# ============================================================
+
+@router.get("/sessions/{session_id}/metadata", response_model=SessionMetadataResponse)
+async def get_metadata(
+    session_id: str,
+    session: OptimizationSession = Depends(get_session)
+):
+    """
+    Get session metadata.
+    
+    Returns the session's user-friendly name, description, tags, and timestamps.
+    """
+    return SessionMetadataResponse(
+        session_id=session.metadata.session_id,
+        name=session.metadata.name,
+        created_at=session.metadata.created_at,
+        last_modified=session.metadata.last_modified,
+        description=session.metadata.description,
+        tags=session.metadata.tags
+    )
+
+
+@router.patch("/sessions/{session_id}/metadata", response_model=SessionMetadataResponse)
+async def update_metadata(
+    session_id: str,
+    request: UpdateMetadataRequest,
+    session: OptimizationSession = Depends(get_session)
+):
+    """
+    Update session metadata.
+    
+    Update the session's name, description, and/or tags. Only provided fields
+    will be updated; omitted fields remain unchanged.
+    """
+    session.update_metadata(
+        name=request.name,
+        description=request.description,
+        tags=request.tags
+    )
+    
+    return SessionMetadataResponse(
+        session_id=session.metadata.session_id,
+        name=session.metadata.name,
+        created_at=session.metadata.created_at,
+        last_modified=session.metadata.last_modified,
+        description=session.metadata.description,
+        tags=session.metadata.tags
+    )
+
+
+# ============================================================
+# Audit Log Endpoints
+# ============================================================
+
+@router.get("/sessions/{session_id}/audit", response_model=AuditLogResponse)
+async def get_audit_log(
+    session_id: str,
+    entry_type: str = None,
+    session: OptimizationSession = Depends(get_session)
+):
+    """
+    Get audit log entries.
+    
+    Retrieves the complete audit trail or filters by entry type.
+    
+    Args:
+        session_id: Session identifier
+        entry_type: Optional filter ('data_locked', 'model_locked', 'acquisition_locked')
+    """
+    if entry_type:
+        entries = session.audit_log.get_entries(entry_type)
+    else:
+        entries = session.audit_log.get_entries()
+    
+    return AuditLogResponse(
+        entries=[AuditEntryResponse(**e.to_dict()) for e in entries],
+        n_entries=len(entries)
+    )
+
+
+@router.post("/sessions/{session_id}/audit/lock", response_model=LockDecisionResponse)
+async def lock_decision(
+    session_id: str,
+    request: LockDecisionRequest,
+    session: OptimizationSession = Depends(get_session)
+):
+    """
+    Lock in a decision to the audit log.
+    
+    Creates an immutable audit entry for data, model, or acquisition decisions.
+    This should be called when the user is satisfied with their configuration
+    and ready to commit the decision to the audit trail.
+    
+    Args:
+        session_id: Session identifier
+        request: Lock decision request
+    """
+    try:
+        if request.lock_type == "data":
+            entry = session.lock_data(notes=request.notes or "")
+            message = "Data decision locked successfully"
+            
+        elif request.lock_type == "model":
+            entry = session.lock_model(notes=request.notes or "")
+            message = "Model decision locked successfully"
+            
+        elif request.lock_type == "acquisition":
+            if not request.strategy or not request.parameters or not request.suggestions:
+                raise HTTPException(
+                    status_code=status.HTTP_400_BAD_REQUEST,
+                    detail="Acquisition lock requires strategy, parameters, and suggestions"
+                )
+            entry = session.lock_acquisition(
+                strategy=request.strategy,
+                parameters=request.parameters,
+                suggestions=request.suggestions,
+                notes=request.notes or ""
+            )
+            message = "Acquisition decision locked successfully"
+            
+        else:
+            raise HTTPException(
+                status_code=status.HTTP_400_BAD_REQUEST,
+                detail=f"Invalid lock_type: {request.lock_type}"
+            )
+        
+        return LockDecisionResponse(
+            success=True,
+            entry=AuditEntryResponse(**entry.to_dict()),
+            message=message
+        )
+        
+    except ValueError as e:
+        raise HTTPException(
+            status_code=status.HTTP_400_BAD_REQUEST,
+            detail=str(e)
+        )
+
+
+@router.get("/sessions/{session_id}/audit/export")
+async def export_audit_markdown(
+    session_id: str,
+    session: OptimizationSession = Depends(get_session)
+):
+    """
+    Export audit log as markdown.
+    
+    Returns the audit trail formatted as markdown for publication methods sections.
+    """
+    markdown = session.export_audit_markdown()
+    
+    return Response(
+        content=markdown,
+        media_type="text/markdown",
+        headers={
+            "Content-Disposition": f"attachment; filename=audit_log_{session_id}.md"
+        }
+    )
+
+
+# ============================================================
+# Session File Management (JSON Format)
+# ============================================================
+
+@router.get("/sessions/{session_id}/download")
+async def download_session(
+    session_id: str,
+    session: OptimizationSession = Depends(get_session)
+):
+    """
+    Download session as JSON file.
+    
+    Downloads the complete session state as a .json file with user-friendly
+    naming support. The file includes metadata, audit log, search space,
+    experiments, and configuration.
+    """
+    # Create temporary file
+    with tempfile.NamedTemporaryFile(mode='w', suffix='.json', delete=False) as f:
+        temp_path = f.name
+        
+    try:
+        # Save session to temp file
+        session.save_session(temp_path)
+        
+        # Use session name for filename (sanitized)
+        filename = session.metadata.name.replace(" ", "_").replace("/", "_")
+        filename = f"{filename}.json"
+        
+        return FileResponse(
+            path=temp_path,
+            media_type="application/json",
+            filename=filename,
+            headers={"Content-Disposition": f"attachment; filename={filename}"}
+        )
+    except Exception as e:
+        # Clean up temp file on error
+        Path(temp_path).unlink(missing_ok=True)
+        raise HTTPException(
+            status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
+            detail=f"Failed to export session: {str(e)}"
+        )
+
+
+@router.post("/sessions/upload", response_model=SessionCreateResponse, status_code=status.HTTP_201_CREATED)
+async def upload_session(file: UploadFile = File(...)):
+    """
+    Upload and restore a session from JSON file.
+    
+    Uploads a .json session file and creates a new session with the restored data.
+    A new session ID will be generated for API use, but the original session ID
+    is preserved in the metadata.
+    """
+    try:
+        # Save uploaded file to temp location
+        with tempfile.NamedTemporaryFile(mode='wb', suffix='.json', delete=False) as f:
+            content = await file.read()
+            f.write(content)
+            temp_path = f.name
+        
+        try:
+            # Load session from file without retraining
+            loaded_session = OptimizationSession.load_session(temp_path, retrain_on_load=False)
+
+            # Create new session in store
+            new_session_id = session_store.create()
+
+            # Replace the session object with loaded one and align metadata
+            try:
+                loaded_session.metadata.session_id = new_session_id
+            except Exception:
+                pass
+
+            session_store._sessions[new_session_id]["session"] = loaded_session
+
+            # Update last accessed
+            session_store._sessions[new_session_id]["last_accessed"] = datetime.now()
+
+            # Persist to disk
+            session_store._save_to_disk(new_session_id)
+            
+            session_info = session_store.get_info(new_session_id)
+            
+            return SessionCreateResponse(
+                session_id=new_session_id,
+                created_at=session_info["created_at"],
+                expires_at=session_info["expires_at"]
+            )
+            
+        finally:
+            # Clean up temp file
+            Path(temp_path).unlink(missing_ok=True)
+            
+    except Exception as e:
+        logger.error(f"Failed to upload session: {e}")
+        raise HTTPException(
+            status_code=status.HTTP_400_BAD_REQUEST,
+            detail=f"Failed to upload session: {str(e)}"
+        )
+
+
+# ============================================================
+# Session Locking Endpoints
+# ============================================================
+
+@router.post("/sessions/{session_id}/lock", response_model=SessionLockResponse)
+async def lock_session(
+    session_id: str,
+    request: SessionLockRequest
+):
+    """
+    Lock a session for external programmatic control.
+    
+    When locked, the web UI should enter monitor-only mode.
+    Returns a lock_token that must be used to unlock.
+    """
+    try:
+        result = session_store.lock_session(
+            session_id=session_id,
+            locked_by=request.locked_by,
+            client_id=request.client_id
+        )
+        
+        # Broadcast lock event to WebSocket clients
+        await broadcast_to_session(session_id, {
+            "event": "lock_status_changed",
+            "locked": True,
+            "locked_by": request.locked_by,
+            "locked_at": result["locked_at"]
+        })
+        
+        return SessionLockResponse(**result)
+    except KeyError:
+        raise HTTPException(
+            status_code=status.HTTP_404_NOT_FOUND,
+            detail=f"Session {session_id} not found or expired"
+        )
+    except Exception as e:
+        raise HTTPException(
+            status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
+            detail=f"Failed to lock session: {str(e)}"
+        )
+
+
+@router.delete("/sessions/{session_id}/lock", response_model=SessionLockResponse)
+async def unlock_session(
+    session_id: str,
+    lock_token: Optional[str] = None
+):
+    """
+    Unlock a session.
+    
+    Optionally provide lock_token for verification.
+    If no token provided, forcibly unlocks (use with caution).
+    """
+    try:
+        result = session_store.unlock_session(session_id=session_id, lock_token=lock_token)
+        
+        # Broadcast unlock event to WebSocket clients
+        await broadcast_to_session(session_id, {
+            "event": "lock_status_changed",
+            "locked": False,
+            "locked_by": None,
+            "locked_at": None
+        })
+        
+        return SessionLockResponse(**result)
+    except KeyError:
+        raise HTTPException(
+            status_code=status.HTTP_404_NOT_FOUND,
+            detail=f"Session {session_id} not found or expired"
+        )
+    except ValueError as e:
+        raise HTTPException(
+            status_code=status.HTTP_403_FORBIDDEN,
+            detail=str(e)
+        )
+    except Exception as e:
+        raise HTTPException(
+            status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
+            detail=f"Failed to unlock session: {str(e)}"
+        )
+
+
+@router.get("/sessions/{session_id}/lock", response_model=SessionLockResponse)
+async def get_lock_status(session_id: str):
+    """
+    Get current lock status of a session.
+    
+    Used by web UI to detect when external controller has taken control
+    and automatically enter monitor mode.
+    """
+    try:
+        result = session_store.get_lock_status(session_id=session_id)
+        return SessionLockResponse(**result)
+    except KeyError:
+        raise HTTPException(
+            status_code=status.HTTP_404_NOT_FOUND,
+            detail=f"Session {session_id} not found or expired"
+        )
+    except Exception as e:
+        raise HTTPException(
+            status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
+            detail=f"Failed to get lock status: {str(e)}"
+        )

api/routers/visualizations.py

@@ -192,12 +192,17 @@ async def get_contour_data(
     grid_df = pd.DataFrame(grid_points)
     
     # CRITICAL FIX: Reorder columns to match training data
-    # The model was trained with a specific column order, we must match it
+    # 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).
     train_data = session.experiment_manager.get_data()
-    train_columns = [col for col in train_data.columns if col != 'Output']
-    
-    # Reorder grid_df to match training column order
-    grid_df = grid_df[train_columns]
+    metadata_cols = {'Iteration', 'Reason', 'Output', 'Noise'}
+    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.
+    # Use reindex so missing columns (shouldn't happen) are filled with the
+    # midpoint/defaults the grid already supplies; this avoids KeyError.
+    grid_df = grid_df.reindex(columns=feature_cols)
     
     # IMPORTANT: The model's predict() method handles preprocessing internally
     # (including categorical encoding), so we can pass the raw DataFrame directly

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.info(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.info(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()
+    }