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

mito_ai/__init__.py

@@ -6,12 +6,14 @@ from jupyter_server.utils import url_path_join
 from mito_ai.completions.handlers import CompletionHandler
 from mito_ai.completions.providers import OpenAIProvider
 from mito_ai.app_builder.handlers import AppBuilderHandler
+from mito_ai.streamlit_preview.handlers import StreamlitPreviewHandler
 from mito_ai.log.urls import get_log_urls
 from mito_ai.version_check import VersionCheckHandler
 from mito_ai.db.urls import get_db_urls
 from mito_ai.settings.urls import get_settings_urls
 from mito_ai.rules.urls import get_rules_urls
 from mito_ai.auth.urls import get_auth_urls
+from mito_ai.streamlit_preview.urls import get_streamlit_preview_urls
 try:
     from _version import __version__
 except ImportError:
@@ -58,6 +60,11 @@ def _load_jupyter_server_extension(server_app) -> None: # type: ignore
             AppBuilderHandler,
             {}
         ),
+        (
+            url_path_join(base_url, "mito-ai", "streamlit-preview"),
+            StreamlitPreviewHandler,
+            {}
+        ),
         (
             url_path_join(base_url, "mito-ai", "version-check"),
             VersionCheckHandler,
@@ -69,8 +76,9 @@ def _load_jupyter_server_extension(server_app) -> None: # type: ignore
     handlers.extend(get_db_urls(base_url))  # type: ignore
     handlers.extend(get_settings_urls(base_url))  # type: ignore
     handlers.extend(get_rules_urls(base_url))  # type: ignore
-    handlers.extend(get_log_urls(base_url))  # type: ignore
+    handlers.extend(get_log_urls(base_url, open_ai_provider.key_type))  # type: ignore
     handlers.extend(get_auth_urls(base_url))  # type: ignore
+    handlers.extend(get_streamlit_preview_urls(base_url))  # type: ignore
     
     web_app.add_handlers(host_pattern, handlers)
     server_app.log.info("Loaded the mito_ai server extension")

mito_ai/_version.py

@@ -1,4 +1,4 @@
 # This file is auto-generated by Hatchling. As such, do not:
 #   - modify
 #   - track in version control e.g. be sure to add to .gitignore
-__version__ = VERSION = '0.1.37'
+__version__ = VERSION = '0.1.38'

mito_ai/app_builder/handlers.py

@@ -13,6 +13,7 @@ from mito_ai.utils.websocket_base import BaseWebSocketHandler
 from mito_ai.app_builder.models import (
     BuildAppReply,
     AppBuilderError,
+    BuildAppRequest,
     ErrorMessage,
     MessageType
 )
@@ -74,7 +75,8 @@ class AppBuilderHandler(BaseWebSocketHandler):
             
             if message_type == MessageType.BUILD_APP.value:
                 # Handle build app request
-                await self._handle_build_app(parsed_message)
+                build_app_request = BuildAppRequest(**parsed_message)
+                await self._handle_build_app(build_app_request)
             else:
                 self.log.error(f"Unknown message type: {message_type}")
                 error = AppBuilderError(
@@ -98,25 +100,24 @@ class AppBuilderHandler(BaseWebSocketHandler):
         latency_ms = round((time.time() - start) * 1000)
         self.log.info(f"App builder handler processed in {latency_ms} ms.")
     
-    async def _handle_build_app(self, message: dict) -> None:
+    async def _handle_build_app(self, message: BuildAppRequest) -> None:
         """Handle a build app request.
         
         Args:
             message: The parsed message.
         """
-        message_id = message.get('message_id', '')  # Default to empty string if not present
-        notebook_path = message.get('notebook_path')
-        app_path = message.get('app_path')
-        jwt_token = message.get('jwt_token', '')  # Extract JWT token from request, default to empty string
+        message_id = message.message_id
+        notebook_path = message.notebook_path
+        jwt_token = message.jwt_token
 
         if not message_id:
             self.log.error("Missing message_id in request")
             return
         
-        if not app_path:
+        if not notebook_path:
             error = AppBuilderError(
                 error_type="InvalidRequest",
-                title="Missing 'path' parameter"
+                title="Missing 'notebook_path' parameter"
             )
             self.reply(BuildAppReply(
                 parent_id=message_id,
@@ -126,31 +127,30 @@ class AppBuilderHandler(BaseWebSocketHandler):
             return
 
         # Validate JWT token if provided
-        if jwt_token and jwt_token != 'placeholder-jwt-token':
-            self.log.info(f"Validating JWT token: {jwt_token[:20]}...")
-            is_valid = self._validate_jwt_token(jwt_token)
-            if not is_valid:
-                self.log.error("JWT token validation failed")
-                error = AppBuilderError(
-                    error_type="Unauthorized",
-                    title="Invalid authentication token",
-                    hint="Please sign in again to deploy your app."
-                )
-                self.reply(BuildAppReply(
-                    parent_id=message_id,
-                    url="",
-                    error=error
-                ))
-                return
-            else:
-                self.log.info("JWT token validation successful")
+        token_preview = jwt_token[:20] if jwt_token else "No token provided"
+        self.log.info(f"Validating JWT token: {token_preview}...")
+        is_valid = self._validate_jwt_token(jwt_token) if jwt_token else False
+        if not is_valid or not jwt_token:
+            self.log.error("JWT token validation failed")
+            error = AppBuilderError(
+                error_type="Unauthorized",
+                title="Invalid authentication token",
+                hint="Please sign in again to deploy your app."
+            )
+            self.reply(BuildAppReply(
+                parent_id=message_id,
+                url="",
+                error=error
+            ))
+            return
         else:
-            self.log.warning("No JWT token provided or using placeholder token")
-
+            self.log.info("JWT token validation successful")
+            
         try:
 
-            success_flag, result_message = await streamlit_handler(str(notebook_path) if notebook_path else "", app_path)
-            if not success_flag:
+            notebook_path = str(notebook_path) if notebook_path else ""
+            success_flag, app_path, result_message = await streamlit_handler(notebook_path)
+            if not success_flag or app_path is None:
                 raise Exception(result_message)
 
             deploy_url = await self._deploy_app(app_path, jwt_token)

mito_ai/app_builder/models.py

@@ -65,7 +65,7 @@ class BuildAppRequest:
     message_id: str
     
     # Path to the app file.
-    path: str
+    notebook_path: str
     
     # JWT token for authorization.
     jwt_token: Optional[str] = None

mito_ai/log/handlers.py

@@ -3,16 +3,22 @@
 
 from dataclasses import dataclass
 import json
-from typing import Any, Final
+from typing import Any, Final, Literal
 import tornado
 import os
 from jupyter_server.base.handlers import APIHandler
-from mito_ai.utils.telemetry_utils import log
+from mito_ai.utils.telemetry_utils import MITO_SERVER_KEY, USER_KEY, log
 
 
 class LogHandler(APIHandler):
     """Handler for logging"""
     
+    def initialize(self, key_type: Literal['mito_server_key', 'user_key']) -> None:
+        """Initialize the log handler"""
+        
+        # The key_type is required so that we know if we can log pro users
+        self.key_type = key_type
+        
     @tornado.web.authenticated
     def put(self) -> None:
         """Log an event"""
@@ -26,6 +32,7 @@ class LogHandler(APIHandler):
         log_event = data['log_event']
         params = data.get('params', {})
         
-        log(log_event, params)
+        key_type = MITO_SERVER_KEY if self.key_type == "mito_server_key" else USER_KEY
+        log(log_event, params, key_type=key_type)
 
 

mito_ai/log/urls.py

@@ -5,7 +5,7 @@ from typing import Any, List, Tuple
 from jupyter_server.utils import url_path_join
 from mito_ai.log.handlers import LogHandler
 
-def get_log_urls(base_url: str) -> List[Tuple[str, Any, dict]]:
+def get_log_urls(base_url: str, key_type: str) -> List[Tuple[str, Any, dict]]:
     """Get all log related URL patterns.
 
     Args:
@@ -15,7 +15,7 @@ def get_log_urls(base_url: str) -> List[Tuple[str, Any, dict]]:
         List of (url_pattern, handler_class, handler_kwargs) tuples
     """
     BASE_URL = base_url + "/mito-ai"
-    
+        
     return [
-        (url_path_join(BASE_URL, "log"), LogHandler, {}),
+        (url_path_join(BASE_URL, "log"), LogHandler, {"key_type": key_type}),
     ]

mito_ai/streamlit_conversion/streamlit_agent_handler.py

@@ -2,8 +2,9 @@
 # Distributed under the terms of the GNU Affero General Public License v3.0 License.
 
 import logging
+import os
 from anthropic.types import MessageParam
-from typing import List, Tuple, cast
+from typing import List, Optional, Tuple, cast
 
 from mito_ai.logger import get_logger
 from mito_ai.streamlit_conversion.streamlit_system_prompt import streamlit_system_prompt
@@ -11,6 +12,7 @@ from mito_ai.streamlit_conversion.validate_and_run_streamlit_code import streaml
 from mito_ai.streamlit_conversion.streamlit_utils import extract_code_blocks, create_app_file, parse_jupyter_notebook_to_extract_required_content
 from mito_ai.utils.anthropic_utils import stream_anthropic_completion_from_mito_server
 from mito_ai.completions.models import MessageType
+from mito_ai.utils.telemetry_utils import log_streamlit_app_creation_error, log_streamlit_app_creation_retry, log_streamlit_app_creation_success
 
 STREAMLIT_AI_MODEL = "claude-3-5-haiku-latest"
 
@@ -93,20 +95,34 @@ class StreamlitCodeGeneration:
         return converted_code
 
 
-async def streamlit_handler(notebook_path: str, app_path: str) -> Tuple[bool, str]:
+async def streamlit_handler(notebook_path: str) -> Tuple[bool, Optional[str], str]:
     """Handler function for streamlit code generation and validation"""
     notebook_code = parse_jupyter_notebook_to_extract_required_content(notebook_path)
     streamlit_code_generator = StreamlitCodeGeneration(notebook_code)
     streamlit_code = await streamlit_code_generator.generate_streamlit_code()
     has_validation_error, error = streamlit_code_validator(streamlit_code)
+    
+    
     tries = 0
     while has_validation_error and tries < 5:
         streamlit_code = await streamlit_code_generator.correct_error_in_generation(error)
         has_validation_error, error = streamlit_code_validator(streamlit_code)
+        
+        if has_validation_error:
+            # TODO: We can't easily get the key type here, so for the beta release
+            # we are just defaulting to the mito server key since that is by far the most common.
+            log_streamlit_app_creation_retry('mito_server_key', MessageType.STREAMLIT_CONVERSION, error)
         tries+=1
 
     if has_validation_error:
-        return False, "Error generating streamlit code by agent"
-
-    success_flag, message = create_app_file(app_path, streamlit_code)
-    return success_flag, message
+        log_streamlit_app_creation_error('mito_server_key', MessageType.STREAMLIT_CONVERSION, error)
+        return False, None, "Error generating streamlit code by agent"
+    
+    app_directory = os.path.dirname(notebook_path)
+    success_flag, app_path, message = create_app_file(app_directory, streamlit_code)
+    
+    if not success_flag:
+        log_streamlit_app_creation_error('mito_server_key', MessageType.STREAMLIT_CONVERSION, message)
+    
+    log_streamlit_app_creation_success('mito_server_key', MessageType.STREAMLIT_CONVERSION)
+    return success_flag, app_path, message

mito_ai/streamlit_conversion/streamlit_system_prompt.py

@@ -28,6 +28,17 @@ CODE STRUCTURE:
 - Handle data loading and processing
 - Organize content with clear sections and headers
 - Include comments explaining key sections
+- Always include the following code at the top of the file so the user does not use the wrong deploy button
+```python
+st.markdown(\"\"\"
+    <style>
+        #MainMenu {visibility: hidden;}
+        .stAppDeployButton {display:none;}
+        footer {visibility: hidden;}
+        .stMainBlockContainer {padding: 2rem 1rem 2rem 1rem;}
+    </style>
+\"\"\", unsafe_allow_html=True)
+```
 
 OUTPUT FORMAT:
 - Provide the complete app.py file code

mito_ai/streamlit_conversion/streamlit_utils.py

@@ -3,7 +3,8 @@
 
 import re
 import json
-from typing import Dict, Tuple, Any
+import os
+from typing import Dict, Optional, Tuple, Any
 
 def extract_code_blocks(message_content: str) -> str:
     """
@@ -27,7 +28,7 @@ def extract_code_blocks(message_content: str) -> str:
     return '\n'.join(matches)
 
 
-def create_app_file(file_path: str, code: str) -> Tuple[bool, str]:
+def create_app_file(app_directory: str, code: str) -> Tuple[bool, Optional[str], str]:
     """
     Create app.py file and write code to it with error handling
 
@@ -40,13 +41,14 @@ def create_app_file(file_path: str, code: str) -> Tuple[bool, str]:
 
     """
     try:
-        with open(file_path+"/app.py", 'w') as f:
+        app_path = os.path.join(app_directory, "app.py")
+        with open(app_path, 'w') as f:
             f.write(code)
-        return True, f"Successfully created {file_path}"
+        return True, app_path, f"Successfully created {app_directory}"
     except IOError as e:
-        return False, f"Error creating file: {str(e)}"
+        return False, None, f"Error creating file: {str(e)}"
     except Exception as e:
-        return False, f"Unexpected error: {str(e)}"
+        return False, None, f"Unexpected error: {str(e)}"
 
 
 def parse_jupyter_notebook_to_extract_required_content(notebook_path: str) -> Dict[str, Any]:

mito_ai/streamlit_conversion/validate_and_run_streamlit_code.py

@@ -42,6 +42,7 @@ class StreamlitValidator:
         self.temp_dir = tempfile.mkdtemp()
         if self.temp_dir is None:
             raise RuntimeError("Failed to create temporary directory")
+        
         app_path = os.path.join(self.temp_dir, "app.py")
 
         with open(app_path, 'w') as f:

mito_ai/streamlit_preview/__init__.py

@@ -0,0 +1,7 @@
+# Copyright (c) Saga Inc.
+# Distributed under the terms of the GNU Affero General Public License v3.0 License.
+
+from .manager import get_preview_manager
+from .handlers import StreamlitPreviewHandler
+
+__all__ = ['get_preview_manager', 'StreamlitPreviewHandler'] 

mito_ai/streamlit_preview/handlers.py

@@ -0,0 +1,161 @@
+# 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 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
+            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)
+            self.finish({"error": f'Internal server 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