mito-ai 0.1.37__py3-none-any.whl → 0.1.39__py3-none-any.whl

mito_ai/streamlit_preview/handlers.py

@@ -0,0 +1,164 @@
+# Copyright (c) Saga Inc.
+# Distributed under the terms of the GNU Affero General Public License v3.0 License.
+
+import os
+import tempfile
+import uuid
+import tornado
+from jupyter_server.base.handlers import APIHandler
+from mito_ai.streamlit_conversion.streamlit_agent_handler import streamlit_handler
+from mito_ai.streamlit_preview.manager import get_preview_manager
+from mito_ai.utils.create import initialize_user
+
+
+class StreamlitPreviewHandler(APIHandler):
+    """REST handler for streamlit preview operations."""
+    
+    def initialize(self) -> None:
+        """Initialize the handler."""
+        self.preview_manager = get_preview_manager()
+    
+    def _resolve_notebook_path(self, notebook_path: str) -> str:
+        """
+        Resolve the notebook path to an absolute path that can be found by the backend.
+        
+        This method handles path resolution issues that can occur in different environments:
+        
+        1. **Test Environment**: Playwright tests create temporary directories with complex
+           paths like 'mitoai_ui_tests-app_builde-ab3a5-n-Test-Preview-as-Streamlit-chromium/'
+           that the backend can't directly access.
+        
+        2. **JupyterHub/Cloud Deployments**: In cloud environments, users may have notebooks
+           in subdirectories that aren't immediately accessible from the server root.
+        
+        3. **Docker Containers**: When running in containers, the working directory and
+           file paths may not align with what the frontend reports.
+        
+        4. **Multi-user Environments**: In enterprise deployments, users may have notebooks
+           in user-specific directories that require path resolution.
+        
+        The method tries multiple strategies:
+        1. If the path is already absolute, return it as-is
+        2. Try to resolve relative to the Jupyter server's root directory
+        3. Search recursively through subdirectories for a file with the same name
+        4. Return the original path if not found (will cause a clear error message)
+        
+        Args:
+            notebook_path (str): The notebook path from the frontend (may be relative or absolute)
+            
+        Returns:
+            str: The resolved absolute path to the notebook file
+        """
+        # If the path is already absolute, return it
+        if os.path.isabs(notebook_path):
+            return notebook_path
+        
+        # Get the Jupyter server's root directory
+        server_root = self.settings.get('server_root_dir', os.getcwd())
+        
+        # Try to find the notebook file in the server root
+        resolved_path = os.path.join(server_root, notebook_path)
+        if os.path.exists(resolved_path):
+            return resolved_path
+        
+        # If not found, try to find it in subdirectories
+        # This handles cases where the notebook is in a subdirectory that the frontend
+        # doesn't know about, or where the path structure differs between frontend and backend
+        for root, dirs, files in os.walk(server_root):
+            if os.path.basename(notebook_path) in files:
+                return os.path.join(root, os.path.basename(notebook_path))
+        
+        # If still not found, return the original path (will cause a clear error)
+        # This ensures we get a meaningful error message rather than a generic "file not found"
+        return os.path.join(os.getcwd(), notebook_path)
+
+    @tornado.web.authenticated
+    async def post(self) -> None:
+        """Start a new streamlit preview.
+        
+        Expected JSON body:
+        {
+            "notebook_path": "path/to/notebook.ipynb"
+        }
+        
+        Returns:
+        {
+            "id": "preview_id",
+            "port": 8501,
+            "url": "http://localhost:8501"
+        }
+        """
+        try:
+            # Parse request body
+            body = self.get_json_body()
+            if body is None:
+                self.set_status(400)
+                self.finish({"error": 'Invalid or missing JSON body'})
+                return
+
+            notebook_path = body.get('notebook_path')
+
+            if not notebook_path:
+                self.set_status(400)
+                self.finish({"error": 'Missing notebook_path parameter'})
+                return
+            
+            # Resolve the notebook path to find the actual file
+            resolved_notebook_path = self._resolve_notebook_path(notebook_path)
+            
+            # Generate preview ID
+            preview_id = str(uuid.uuid4())
+            
+            # Generate streamlit code using existing handler
+            print('notebook_path', notebook_path)
+            success, app_path, message = await streamlit_handler(resolved_notebook_path)
+            
+            if not success or app_path is None:
+                self.set_status(500)
+                self.finish({"error": f'Failed to generate streamlit code: {message}'})
+                return
+            
+            # Start streamlit preview
+            resolved_app_directory = os.path.dirname(resolved_notebook_path)
+            success, message, port = self.preview_manager.start_streamlit_preview(resolved_app_directory, preview_id)
+            
+            if not success:
+                self.set_status(500)
+                self.finish({"error": f'Failed to start preview: {message}'})
+                return
+            
+            # Return success response - APIHandler automatically handles JSON serialization
+            self.finish({
+                'id': preview_id,
+                'port': port,
+                'url': f'http://localhost:{port}'
+            })
+                
+        except Exception as e:
+            print(f"Error in streamlit preview handler: {e}")
+            self.set_status(500)
+            
+            # Respond with the error
+            self.finish({"error": str(e)})
+    
+    @tornado.web.authenticated
+    def delete(self, preview_id: str) -> None:
+        """Stop a streamlit preview."""
+        try:
+            if not preview_id:
+                self.set_status(400)
+                self.finish({"error": 'Missing preview_id parameter'})
+                return
+            
+            # Stop the preview
+            stopped = self.preview_manager.stop_preview(preview_id)
+            
+            if stopped:
+                self.set_status(204)  # No content
+            else:
+                self.set_status(404)
+                self.finish({"error": f'Preview {preview_id} not found'})
+                
+        except Exception as e:
+            self.set_status(500)
+            self.finish({"error": f'Internal server error: {str(e)}'})

mito_ai/streamlit_preview/manager.py

@@ -0,0 +1,159 @@
+# Copyright (c) Saga Inc.
+# Distributed under the terms of the GNU Affero General Public License v3.0 License.
+
+import os
+import socket
+import subprocess
+import tempfile
+import time
+import threading
+import requests
+from typing import Dict, Optional, Tuple
+from dataclasses import dataclass
+from mito_ai.logger import get_logger
+
+
+@dataclass
+class PreviewProcess:
+    """Data class to track a streamlit preview process."""
+    proc: subprocess.Popen
+    port: int
+
+
+class StreamlitPreviewManager:
+    """Manages streamlit preview processes and their lifecycle."""
+    
+    def __init__(self) -> None:
+        self._previews: Dict[str, PreviewProcess] = {}
+        self._lock = threading.Lock()
+        self.log = get_logger()
+        
+    def get_free_port(self) -> int:
+        """Get a free port for streamlit to use."""
+        with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
+            s.bind(('', 0))
+            s.listen(1)
+            port = int(s.getsockname()[1])
+        
+        return port
+    
+    def start_streamlit_preview(self, app_directory: str, preview_id: str) -> Tuple[bool, str, Optional[int]]:
+        """Start a streamlit preview process.
+        
+        Args:
+            app_code: The streamlit app code to run
+            preview_id: Unique identifier for this preview
+            
+        Returns:
+            Tuple of (success, message, port)
+        """
+        try:
+            
+            # Get free port
+            port = self.get_free_port()
+            
+            # Start streamlit process
+            cmd = [
+                "streamlit", "run", 'app.py', # Since we run this command from the app_directory, we always just run app.py 
+                "--server.port", str(port),
+                "--server.headless", "true",
+                "--server.address", "localhost",
+                "--server.enableXsrfProtection", "false",
+                "--server.runOnSave", "true",  # auto-reload when app.py is saved
+                "--logger.level", "error"
+            ]
+            
+            # TODO: Security considerations for production:
+            # - Consider enabling XSRF protection if needed, but we might already get this with the APIHandler?
+            # - Add authentication headers to streamlit
+            
+            proc = subprocess.Popen(
+                cmd,
+                stdout=subprocess.PIPE,
+                stderr=subprocess.PIPE,
+                text=True,
+                cwd=app_directory
+            )
+            
+            # Wait for app to be ready
+            ready = self._wait_for_app_ready(port)
+            if not ready:
+                proc.terminate()
+                proc.wait()
+                return False, "Streamlit app failed to start", None
+            
+            # Register the process
+            with self._lock:
+                self._previews[preview_id] = PreviewProcess(
+                    proc=proc,
+                    port=port,
+                )
+            
+            self.log.info(f"Started streamlit preview {preview_id} on port {port}")
+            return True, "Preview started successfully", port
+            
+        except Exception as e:
+            self.log.error(f"Error starting streamlit preview: {e}")
+            return False, f"Failed to start preview: {str(e)}", None
+    
+    def _wait_for_app_ready(self, port: int, timeout: int = 30) -> bool:
+        """Wait for streamlit app to be ready on the given port."""
+        start_time = time.time()
+        
+        while time.time() - start_time < timeout:
+            try:
+                response = requests.get(f"http://localhost:{port}", timeout=5)
+                if response.status_code == 200:
+                    return True
+            except requests.exceptions.RequestException as e:
+                print(f"Error waiting for app to be ready: {e}")
+                pass
+            
+            time.sleep(1)
+        
+        return False
+    
+    def stop_preview(self, preview_id: str) -> bool:
+        """Stop a streamlit preview process.
+        
+        Args:
+            preview_id: The preview ID to stop
+            
+        Returns:
+            True if stopped successfully, False if not found
+        """
+        print(f"Stopping preview {preview_id}")
+        with self._lock:
+            if preview_id not in self._previews:
+                return False
+            
+            preview = self._previews[preview_id]
+            
+            # Terminate process
+            try:
+                preview.proc.terminate()
+                preview.proc.wait(timeout=5)
+            except subprocess.TimeoutExpired:
+                preview.proc.kill()
+                preview.proc.wait()
+            except Exception as e:
+                self.log.error(f"Error terminating process {preview_id}: {e}")
+            
+            # Remove from registry
+            del self._previews[preview_id]
+            
+            self.log.info(f"Stopped streamlit preview {preview_id}")
+            return True
+    
+    def get_preview(self, preview_id: str) -> Optional[PreviewProcess]:
+        """Get a preview process by ID."""
+        with self._lock:
+            return self._previews.get(preview_id)
+
+# Global instance
+_preview_manager = StreamlitPreviewManager()
+
+
+def get_preview_manager() -> StreamlitPreviewManager:
+    """Get the global preview manager instance."""
+    return _preview_manager 

mito_ai/streamlit_preview/urls.py

@@ -0,0 +1,22 @@
+# Copyright (c) Saga Inc.
+# Distributed under the terms of the GNU Affero General Public License v3.0 License.
+
+from typing import Any, List, Tuple
+from jupyter_server.utils import url_path_join
+from mito_ai.streamlit_preview.handlers import StreamlitPreviewHandler
+
+def get_streamlit_preview_urls(base_url: str) -> List[Tuple[str, Any, dict]]:
+    """Get all streamlit preview related URL patterns.
+
+    Args:
+        base_url: The base URL for the Jupyter server
+
+    Returns:
+        List of (url_pattern, handler_class, handler_kwargs) tuples
+    """
+    BASE_URL = base_url + "/mito-ai"
+    
+    return [
+        (url_path_join(BASE_URL, "streamlit-preview"), StreamlitPreviewHandler, {}),
+        (url_path_join(BASE_URL, "streamlit-preview/(.+)"), StreamlitPreviewHandler, {}),
+    ]