mito-ai 0.1.50__py3-none-any.whl

mito_ai/utils/utils.py

@@ -0,0 +1,89 @@
+# Copyright (c) Saga Inc.
+# Distributed under the terms of the GNU Affero General Public License v3.0 License.
+
+import os
+import pkg_resources
+import subprocess
+import sys
+import uuid
+from typing import List, Optional, Tuple, Union
+from tornado.httpclient import AsyncHTTPClient
+
+def get_random_id() -> str:
+    """
+    Creates a new random ID for the user, which for any given user,
+    should only happen once.
+    """
+    return str(uuid.uuid1())
+
+def is_running_test() -> bool:
+    """
+    A helper function that quickly returns if the current code is running 
+    inside of a test, which is useful for making sure we don't generate 
+    tons of logs.
+    """
+    # Pytest injects PYTEST_CURRENT_TEST into the current environment when running
+    running_pytests = "PYTEST_CURRENT_TEST" in os.environ
+    
+    # Github injects CI into the environment when running
+    running_ci = 'CI' in os.environ and os.environ['CI'] is not None
+
+    return running_pytests or running_ci
+
+def get_installed_packages() -> List[str]:
+    """
+    Get a list of all installed packages.
+    """
+    return [
+        package.key
+        for package in sorted(pkg_resources.working_set, key=lambda x: x.key)
+    ]
+
+def install_packages(packages: List[str]) -> dict[str, Union[bool, str, None]]:
+    """
+    Install a list of packages.
+    
+    Returns:
+        dict: A dictionary containing:
+            - 'success': bool
+            - 'error': error message captured from the subprocess call or None if no error
+    """
+    result: dict[str, Union[bool, str, None]] = {
+        'success': True,
+        'error': None
+    }
+    
+    for package in packages:
+        try:
+            subprocess.check_call([sys.executable, "-m", "pip", "install", package])
+        except subprocess.CalledProcessError as e:
+            result['success'] = False
+            result['error'] = str(e)
+    
+    return result
+
+def _create_http_client(timeout: int, max_retries: int) -> Tuple[AsyncHTTPClient, Optional[int]]:
+    """
+    Create an HTTP client with appropriate timeout settings.
+    
+    Args:
+        timeout: The timeout in seconds
+        max_retries: The maximum number of retries
+        
+    Returns:
+        A tuple containing the HTTP client and the timeout value in milliseconds
+    """
+    from .utils import is_running_test  # local import to avoid circular import if needed
+    if is_running_test():
+        # If we are running in a test environment, setting the request_timeout fails for some reason.
+        http_client = AsyncHTTPClient(defaults=dict(user_agent="Mito-AI client"))
+        http_client_timeout = None
+    else:
+        # To avoid 599 client timeout errors, we set the request_timeout. By default, the HTTP client 
+        # timesout after 20 seconds. We update this to match the timeout we give to OpenAI. 
+        # The OpenAI timeouts are denoted in seconds, whereas the HTTP client expects milliseconds. 
+        # We also give the HTTP client a 10 second buffer to account for
+        http_client_timeout = timeout * 1000 * max_retries + 10000
+        http_client = AsyncHTTPClient(defaults=dict(user_agent="Mito-AI client", request_timeout=http_client_timeout))
+    
+    return http_client, http_client_timeout

mito_ai/utils/version_utils.py

@@ -0,0 +1,94 @@
+# Copyright (c) Saga Inc.
+# Distributed under the terms of the GNU Affero General Public License v3.0 License.
+
+import os
+from typing import cast
+from .schema import UJ_MITOSHEET_ENTERPRISE, UJ_MITOSHEET_PRO
+from .db import get_user_field
+
+# Check if helper packages are installed
+try:
+    import mitosheet_helper_pro
+    MITOSHEET_HELPER_PRO = True
+except ImportError:
+    MITOSHEET_HELPER_PRO = False
+try:
+    import mitosheet_helper_enterprise  
+    MITOSHEET_HELPER_ENTERPRISE = True
+except ImportError:
+    MITOSHEET_HELPER_ENTERPRISE = False
+
+try:
+    import mitosheet_private 
+    MITOSHEET_PRIVATE = True
+except ImportError:
+    MITOSHEET_PRIVATE = False
+
+# This is a legacy helper that we don't use anymore, however, we're keeping it for now
+# for backwards compatibility, since I'm not 100% confident that nobody is currently using it.
+try:
+    import mitosheet_helper_private 
+    MITOSHEET_HELPER_PRIVATE = True
+except ImportError:
+    MITOSHEET_HELPER_PRIVATE = False
+
+
+def is_pro() -> bool:
+    """
+    Helper function for returning if this is a
+    pro deployment of mito
+    """
+
+    # This package overides the user.json
+    if MITOSHEET_HELPER_PRO:
+        return True
+
+    # This package overides the user.json
+    if MITOSHEET_PRIVATE:
+        return True
+
+    # Check if the config is set
+    # TODO: Check if the mito config pro is set to true.
+    # I don't think that any user is on pro via this method
+
+    # If you're on Mito Enterprise, then you get all Mito Pro features
+    if is_enterprise():
+        return True
+
+    pro = get_user_field(UJ_MITOSHEET_PRO)
+    if pro is None:
+        return False
+    
+    return bool(pro)
+
+def is_enterprise() -> bool:
+    """
+    Helper function for returning if this is a Mito Enterprise
+    users
+    """
+
+    # This package overides the user.json
+    if MITOSHEET_HELPER_ENTERPRISE:
+        return True
+    
+    # TODO: Check if the mito config enterprise is set to true. 
+    # I don't think that any user is on enterprise via this method
+    
+    is_enterprise = get_user_field(UJ_MITOSHEET_ENTERPRISE)
+    if is_enterprise is None:
+        return False
+
+    # TODO: Check if someone has a temp enterprise license set
+
+    return bool(is_enterprise)
+
+
+def is_mitosheet_private() -> bool:
+    """
+    Helper function for returning if this is a Mito Private
+    users
+    """
+
+    if MITOSHEET_PRIVATE:
+        return True
+    return False

mito_ai/utils/websocket_base.py

@@ -0,0 +1,88 @@
+# Copyright (c) Saga Inc.
+# Distributed under the terms of the GNU Affero General Public License v3.0 License.
+
+import json
+import logging
+from dataclasses import asdict
+from http import HTTPStatus
+from typing import Any, Dict, Optional
+
+import tornado
+from jupyter_core.utils import ensure_async
+from jupyter_server.base.handlers import JupyterHandler
+from tornado.websocket import WebSocketHandler
+
+from mito_ai.logger import get_logger
+
+
+class BaseWebSocketHandler(JupyterHandler, WebSocketHandler):
+    """Base WebSocket handler with common functionality for Mito AI services."""
+
+    def initialize(self) -> None:
+        """Initialize the WebSocket handler."""
+        super().initialize()
+        self.log.debug("Initializing websocket connection %s", self.request.path)
+
+    @property
+    def log(self) -> logging.Logger:
+        """Use Mito AI logger."""
+        return get_logger()
+
+    async def pre_get(self) -> None:
+        """Handle WebSocket authentication/authorization."""
+        # Authenticate the request before opening the WebSocket
+        user = self.current_user
+        if user is None:
+            self.log.warning("Couldn't authenticate WebSocket connection")
+            raise tornado.web.HTTPError(HTTPStatus.UNAUTHORIZED)
+
+        # Authorize the user
+        if not await ensure_async(
+            self.authorizer.is_authorized(self, user, "execute", "mito-ai")
+        ):
+            raise tornado.web.HTTPError(HTTPStatus.FORBIDDEN)
+
+    async def get(self, *args: Any, **kwargs: Dict[str, Any]) -> None:
+        """Get an event to open a socket or check service availability."""
+        # Check if this is just a service availability check
+        if self.get_query_argument('check_availability', None) == 'true':
+            self.set_status(HTTPStatus.OK)
+            self.finish()
+            return
+
+        await ensure_async(self.pre_get())  # type: ignore
+
+        reply = super().get(*args, **kwargs)
+        if reply is not None:
+            await reply
+
+    def reply(self, reply: Any) -> None:
+        """Write a reply object to the WebSocket connection.
+
+        Args:
+            reply: The reply object, which must be a dataclass instance.
+        """
+        message = asdict(reply)
+        super().write_message(message)
+
+    def on_close(self) -> None:
+        """Invoked when the WebSocket is closed."""
+        self.log.debug("WebSocket closed: %s", self.request.path)
+
+    def parse_message(self, message: str) -> Dict[str, Any]:
+        """Parse an incoming message.
+        
+        Args:
+            message: The JSON message string to parse.
+            
+        Returns:
+            The parsed message as a dict.
+            
+        Raises:
+            ValueError: If the message is not valid JSON.
+        """
+        try:
+            return json.loads(message) # type: ignore
+        except ValueError as e:
+            self.log.error("Invalid message: %s", e)
+            raise 

mito_ai/version_check.py

@@ -0,0 +1,60 @@
+# Copyright (c) Saga Inc.
+# Distributed under the terms of the GNU Affero General Public License v3.0 License.
+
+"""
+Version check endpoint for Mito AI to compare local and PyPI versions.
+"""
+import json
+import logging
+import time
+from functools import lru_cache
+from typing import Tuple, Optional, Any
+import requests
+import tornado.web
+from tornado.web import RequestHandler
+from mito_ai._version import __version__
+
+logger = logging.getLogger(__name__)
+
+class VersionCheckHandler(RequestHandler):
+    """Handler for checking the latest version of Mito AI on PyPI."""
+    
+    # Cache PyPI results for 1 hour (3600 seconds)
+    @staticmethod
+    @lru_cache(maxsize=1)
+    def _get_latest_version() -> Tuple[Optional[str], float]:
+        """Get the latest version from PyPI with caching."""
+        cache_time = time.time()
+        try:
+            response = requests.get("https://pypi.org/pypi/mito-ai/json", timeout=3)
+            response.raise_for_status()
+            pypi_data = response.json()
+            latest_version = pypi_data.get("info", {}).get("version", "")
+            return latest_version, cache_time
+        except (requests.RequestException, json.JSONDecodeError) as e:
+            logger.warning(f"Failed to fetch latest version from PyPI: {str(e)}")
+            return None, cache_time
+    
+    def get(self) -> None:
+        """Get the current and latest Mito AI versions."""
+        try:
+            # Get current package version from _version.py
+            current_version = __version__
+            
+            # Get latest version from PyPI with caching
+            latest_version, cache_time = self._get_latest_version()
+            cache_age = time.time() - cache_time
+            
+            # Include cache information in response
+            self.set_header("Content-Type", "application/json")
+            self.write(json.dumps({
+                "current_version": current_version,
+                "latest_version": latest_version,
+                "cache_age_seconds": int(cache_age)
+            }))
+        except Exception as e:
+            logger.error(f"Error in version check: {str(e)}")
+            self.set_status(500)
+            self.write(json.dumps({
+                "error": "Failed to check versions"
+            })) 

mito_ai-0.1.50.data/data/etc/jupyter/jupyter_server_config.d/mito_ai.json

@@ -0,0 +1,7 @@
+{
+  "ServerApp": {
+    "jpserver_extensions": {
+      "mito_ai": true
+    }
+  }
+}