mito-ai 0.1.43__py3-none-any.whl → 0.1.44__py3-none-any.whl

mito_ai/__init__.py

@@ -5,7 +5,7 @@ from typing import List, Dict
 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.app_deploy.handlers import AppDeployHandler
 from mito_ai.streamlit_preview.handlers import StreamlitPreviewHandler
 from mito_ai.log.urls import get_log_urls
 from mito_ai.version_check import VersionCheckHandler
@@ -71,8 +71,8 @@ def _load_jupyter_server_extension(server_app) -> None: # type: ignore
             {"llm": open_ai_provider},
         ),
         (
-            url_path_join(base_url, "mito-ai", "app-builder"),
-            AppBuilderHandler,
+            url_path_join(base_url, "mito-ai", "app-deploy"),
+            AppDeployHandler,
             {}
         ),
         (

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.43'
+__version__ = VERSION = '0.1.44'

mito_ai/{app_builder → app_deploy}/__init__.py

@@ -3,4 +3,4 @@
 
 """App builder module for Mito AI."""
 
-from .handlers import AppBuilderHandler 
+from .handlers import AppDeployHandler 

mito_ai/{app_builder → app_deploy}/handlers.py

@@ -7,13 +7,14 @@ import logging
 from typing import Any, Union, Optional
 import zipfile
 import tempfile
+from mito_ai.streamlit_conversion.streamlit_utils import get_app_path
 from mito_ai.utils.create import initialize_user
 from mito_ai.utils.version_utils import is_pro
 from mito_ai.utils.websocket_base import BaseWebSocketHandler
-from mito_ai.app_builder.models import (
-    BuildAppReply,
-    AppBuilderError,
-    BuildAppRequest,
+from mito_ai.app_deploy.models import (
+    DeployAppReply,
+    AppDeployError,
+    DeployAppRequest,
     ErrorMessage,
     MessageType
 )
@@ -23,8 +24,8 @@ from mito_ai.constants import ACTIVE_STREAMLIT_BASE_URL
 import requests
 
 
-class AppBuilderHandler(BaseWebSocketHandler):
-    """Handler for app building requests."""
+class AppDeployHandler(BaseWebSocketHandler):
+    """Handler for app deploy requests."""
     
     def initialize(self) -> None:
         """Initialize the WebSocket handler."""
@@ -57,6 +58,7 @@ class AppBuilderHandler(BaseWebSocketHandler):
         Args:
             message: The message received on the WebSocket.
         """
+        
         start = time.time()
         
         # Convert bytes to string if needed
@@ -73,13 +75,13 @@ class AppBuilderHandler(BaseWebSocketHandler):
             parsed_message = self.parse_message(message)
             message_type = parsed_message.get('type')
             
-            if message_type == MessageType.BUILD_APP.value:
+            if message_type == MessageType.DEPLOY_APP.value:
                 # Handle build app request
-                build_app_request = BuildAppRequest(**parsed_message)
-                await self._handle_build_app(build_app_request)
+                deploy_app_request = DeployAppRequest(**parsed_message)
+                await self._handle_deploy_app(deploy_app_request)
             else:
                 self.log.error(f"Unknown message type: {message_type}")
-                error = AppBuilderError(
+                error = AppDeployError(
                     error_type="InvalidRequest",
                     title=f"Unknown message type: {message_type}"
                 )
@@ -87,11 +89,11 @@ class AppBuilderHandler(BaseWebSocketHandler):
                 
         except ValueError as e:
             self.log.error("Invalid app builder request", exc_info=e)
-            error = AppBuilderError.from_exception(e)
+            error = AppDeployError.from_exception(e)
             self.reply(ErrorMessage(**error.__dict__))
         except Exception as e:
             self.log.error("Error handling app builder message", exc_info=e)
-            error = AppBuilderError.from_exception(
+            error = AppDeployError.from_exception(
                 e, 
                 hint="An error occurred while building the app. Please check the logs for details."
             )
@@ -100,7 +102,7 @@ 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: BuildAppRequest) -> None:
+    async def _handle_deploy_app(self, message: DeployAppRequest) -> None:
         """Handle a build app request.
         
         Args:
@@ -109,17 +111,17 @@ class AppBuilderHandler(BaseWebSocketHandler):
         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 notebook_path:
-            error = AppBuilderError(
+            error = AppDeployError(
                 error_type="InvalidRequest",
                 title="Missing 'notebook_path' parameter"
             )
-            self.reply(BuildAppReply(
+            self.reply(DeployAppReply(
                 parent_id=message_id,
                 url="",
                 error=error
@@ -132,12 +134,12 @@ class AppBuilderHandler(BaseWebSocketHandler):
         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 = AppDeployError(
                 error_type="Unauthorized",
                 title="Invalid authentication token",
                 hint="Please sign in again to deploy your app."
             )
-            self.reply(BuildAppReply(
+            self.reply(DeployAppReply(
                 parent_id=message_id,
                 url="",
                 error=error
@@ -150,25 +152,34 @@ class AppBuilderHandler(BaseWebSocketHandler):
             notebook_path = str(notebook_path) if notebook_path else ""
 
             app_directory = os.path.dirname(notebook_path)
-            app_path = os.path.join(app_directory, "app.py")
-
-            if not os.path.exists(app_path):
-                success_flag, app_path_result, result_message = await streamlit_handler(notebook_path)
-                if not success_flag or app_path_result is None:
-                    raise Exception(result_message)
-
+            
+            # Check if the app.py file exists
+            app_path = get_app_path(app_directory)
+            if app_path is None:
+                error = AppDeployError(
+                    error_type="AppNotFound",
+                    title="App not found",
+                    hint="Please make sure the app.py file exists in the same directory as the notebook."
+                )
+                self.reply(DeployAppReply(
+                    parent_id=message_id,
+                    url="",
+                    error=error
+                ))
+            
+            # Finally, deploy the app
             deploy_url = await self._deploy_app(app_directory, jwt_token)
 
             # Send the response
-            self.reply(BuildAppReply(
+            self.reply(DeployAppReply(
                 parent_id=message_id,
                 url=deploy_url
             ))
             
         except Exception as e:
             self.log.error(f"Error building app: {e}", exc_info=e)
-            error = AppBuilderError.from_exception(e)
-            self.reply(BuildAppReply(
+            error = AppDeployError.from_exception(e)
+            self.reply(DeployAppReply(
                 parent_id=message_id,
                 url="",
                 error=error

mito_ai/{app_builder → app_deploy}/models.py

@@ -7,13 +7,13 @@ from typing import Literal, Optional
 
 
 class MessageType(str, Enum):
-    """Types of app builder messages."""
-    BUILD_APP = "build-app"
+    """Types of app deploy messages."""
+    DEPLOY_APP = "deploy-app"
 
 
 @dataclass(frozen=True)
-class AppBuilderError:
-    """Error information for app builder operations."""
+class AppDeployError:
+    """Error information for app deploy operations."""
     
     # Error type.
     error_type: str
@@ -28,7 +28,7 @@ class AppBuilderError:
     hint: Optional[str] = None
     
     @classmethod
-    def from_exception(cls, e: Exception, hint: Optional[str] = None) -> "AppBuilderError":
+    def from_exception(cls, e: Exception, hint: Optional[str] = None) -> "AppDeployError":
         """Create an error from an exception.
 
         Args:
@@ -47,7 +47,7 @@ class AppBuilderError:
 
 
 @dataclass(frozen=True)
-class ErrorMessage(AppBuilderError):
+class ErrorMessage(AppDeployError):
     """Error message."""
     
     # Message type.
@@ -55,11 +55,11 @@ class ErrorMessage(AppBuilderError):
 
 
 @dataclass(frozen=True)
-class BuildAppRequest:
-    """Request to build an app."""
+class DeployAppRequest:
+    """Request to deploy an app."""
     
     # Request type.
-    type: Literal["build-app"]
+    type: Literal["deploy-app"]
     
     # Message ID.
     message_id: str
@@ -72,8 +72,8 @@ class BuildAppRequest:
 
 
 @dataclass(frozen=True)
-class BuildAppReply:
-    """Reply to a build app request."""
+class DeployAppReply:
+    """Reply to a deplpy app request."""
     
     # ID of the request message this is replying to.
     parent_id: str
@@ -82,7 +82,7 @@ class BuildAppReply:
     url: str
     
     # Optional error information.
-    error: Optional[AppBuilderError] = None
+    error: Optional[AppDeployError] = None
     
     # Type of reply.
-    type: Literal["build-app"] = "build-app" 
+    type: Literal["deploy-app"] = "deploy-app" 

mito_ai/app_manager/handlers.py

@@ -12,6 +12,8 @@ from mito_ai.app_manager.models import (
     AppManagerError,
     ManageAppRequest,
     ManageAppReply,
+    CheckAppStatusRequest,
+    CheckAppStatusReply,
     ErrorMessage,
     MessageType
 )
@@ -57,6 +59,10 @@ class AppManagerHandler(BaseWebSocketHandler):
                 # Handle manage app request
                 manage_app_request = ManageAppRequest(**parsed_message)
                 await self._handle_manage_app(manage_app_request)
+            elif message_type == MessageType.CHECK_APP_STATUS.value:
+                # Handle check app status request
+                check_status_request = CheckAppStatusRequest(**parsed_message)
+                await self._handle_check_app_status(check_status_request)
             else:
                 self.log.error(f"Unknown message type: {message_type}")
                 error_response = ErrorMessage(
@@ -131,4 +137,31 @@ class AppManagerHandler(BaseWebSocketHandler):
                 error=error,
                 message_id=request.message_id
             )
+            self.reply(error_reply)
+
+    async def _handle_check_app_status(self, request: CheckAppStatusRequest) -> None:
+        """Handle a check app status request."""
+        self.log.info("In check app status")
+        try:
+            # Make a HEAD request to check if the app URL is accessible
+            response = requests.head(request.app_url, timeout=10, verify=False)
+            self.log.debug(f"Is app accessible: {response.status_code}")
+            is_accessible = response.status_code==200
+
+            # Create successful response
+            reply = CheckAppStatusReply(
+                is_accessible=is_accessible
+            )
+
+            self.reply(reply)
+
+        except Exception as e:
+            self.log.error(f"Error checking app status: {e}", exc_info=e)
+            error = AppManagerError.from_exception(e)
+
+            # Return error response
+            error_reply = CheckAppStatusReply(
+                is_accessible=False,
+                error=error
+            )
             self.reply(error_reply)

mito_ai/app_manager/models.py

@@ -8,6 +8,7 @@ from typing import List, Optional
 class MessageType(str, Enum):
     """Types of app manager messages."""
     MANAGE_APP = "manage-app"
+    CHECK_APP_STATUS = "check-app-status"
 
 
 @dataclass(frozen=True)
@@ -15,7 +16,7 @@ class ManageAppRequest:
     """Request to manage apps."""
     type: str = "manage-app"
     jwt_token: Optional[str] = None
-    message_id: Optional[str] = None  
+    message_id: Optional[str] = None
 
 @dataclass(frozen=True)
 class App:
@@ -48,6 +49,19 @@ class ManageAppReply:
     error: Optional[AppManagerError] = None
     message_id: Optional[str] = None  
 
+@dataclass(frozen=True)
+class CheckAppStatusRequest:
+    """Request to check app status."""
+    app_url: str
+    type: str = "check-app-status"
+
+@dataclass(frozen=True)
+class CheckAppStatusReply:
+    """Reply to a check app status request."""
+    is_accessible: bool
+    type: str = "check-app-status"
+    error: Optional[AppManagerError] = None
+
 @dataclass(frozen=True)
 class ErrorMessage:
     """Error message."""

mito_ai/completions/handlers.py

@@ -67,6 +67,7 @@ class CompletionHandler(JupyterHandler, WebSocketHandler):
         self._llm = llm
         self.is_pro = is_pro()
         self._selected_model = FALLBACK_MODEL
+        self.is_electron = False
         identify(llm.key_type)
         
     @property
@@ -128,6 +129,18 @@ class CompletionHandler(JupyterHandler, WebSocketHandler):
             parsed_message = json.loads(message)
             metadata_dict = parsed_message.get('metadata', {})
             type: MessageType = MessageType(parsed_message.get('type'))
+            
+            # Extract environment information from the message
+            environment = parsed_message.get('environment', {})
+            if environment:
+                is_electron = environment.get('isElectron', None)
+                if is_electron is not None:
+                    if is_electron != self.is_electron:
+                        # If the is_electron status is different, log it
+                        identify(key_type=self._llm.key_type, is_electron=is_electron)
+                    
+                    self.is_electron = is_electron
+
         except ValueError as e:
             self.log.error("Invalid completion request.", exc_info=e)
             return

mito_ai/completions/models.py

@@ -3,7 +3,7 @@
 
 import traceback
 from dataclasses import dataclass, field
-from typing import List, Literal, Optional, NewType, Dict
+from typing import List, Literal, Optional, NewType, Dict, Any
 from openai.types.chat import ChatCompletionMessageParam
 from enum import Enum
 from pydantic import BaseModel
@@ -155,6 +155,9 @@ class CompletionRequest:
     # Whether to stream the response (if supported by the model).
     stream: bool = False
     
+    # Environment information from the client
+    environment: Optional[Dict[str, Any]] = None
+    
     
 @dataclass(frozen=True)
 class AICapabilities:

mito_ai/completions/prompt_builders/agent_system_message.py

@@ -51,7 +51,8 @@ Format:
         code: str
         code_summary: str
         cell_type: 'code' | 'markdown'
-    }}
+    }},
+    analysis_assumptions: Optional[List[str]]
 }}
 
 Important information:
@@ -60,7 +61,7 @@ Important information:
 3. The code should be the full contents of that updated code cell. The code that you return will overwrite the existing contents of the code cell so it must contain all necessary code.
 4. The code_summary must be a very short phrase (1–5 words maximum) that begins with a verb ending in "-ing" (e.g., "Loading data", "Filtering rows", "Calculating average", "Plotting revenue"). Avoid full sentences or explanations—this should read like a quick commit message or code label, not a description.
 5. Important: Only use the CELL_UPDATE tool if you want to add/modify a notebook cell in response to the user's request. If the user is just sending you a friendly greeting or asking you a question about yourself, you SHOULD NOT USE A CELL_UPDATE tool because it does not require modifying the notebook. Instead, just use the FINISHED_TASK response.
-6. The assumptions is an optional list of critical assumptions that you made about the data or analysis approach. The assumptions you list here will be displayed to the user so that they can confirm or correct the assumptions. For example: ["NaN values in the impressions column represent 0 impressions", "Only crashes with pedestrian or cyclist fatalities are considered fatal crashes", "Intervention priority combines both volume and severity to identify maximum impact opportunities"].
+6. The analysis_assumptions is an optional list of critical assumptions that you made about the data or analysis approach. The assumptions you list here will be displayed to the user so that they can confirm or correct the assumptions. For example: ["NaN values in the impressions column represent 0 impressions", "Only crashes with pedestrian or cyclist fatalities are considered fatal crashes", "Intervention priority combines both volume and severity to identify maximum impact opportunities"].
 7. Only include important data and analytical assumptions that if incorrect would fundamentally change your analysis conclusions. These should be data handling decisions, methodological choices, and definitional boundaries. Do not include: obvious statements ("Each record is counted once"), result interpretation guidance ("Gaps in the plot represent zero values"), display choices ("Data is sorted for clarity"), internal reasoning ("Bar chart is better than line plot"), or environment assumptions ("Library X is installed"). Prioritize quality over quantity - include only the most critical assumptions or omit the field entirely if there are no critical assumptions made in this step that have not already be shared with the user. If you ever doubt whether an assumption is critical enough to be shared with the user as an assumption, don't include it. Most messages should not include an assumption. 
 8. Do not include the same assumption or variations of the same assumption multiple times in the same conversation. Once you have presented the assumption to the user, they will already have the opportunity to confirm or correct it so do not include it again.
 
@@ -77,7 +78,8 @@ Format:
         code: str   
         code_summary: str
         cell_type: 'code' | 'markdown'
-    }}
+    }},
+    analysis_assumptions: Optional[List[str]]
 }}
 
 Important information:
@@ -86,7 +88,7 @@ Important information:
 3. The code should be the full contents of that updated code cell. The code that you return will overwrite the existing contents of the code cell so it must contain all necessary code.
 4. code_summary must be a very short phrase (1–5 words maximum) that begins with a verb ending in "-ing" (e.g., "Loading data", "Filtering rows", "Calculating average", "Plotting revenue"). Avoid full sentences or explanations—this should read like a quick commit message or code label, not a description.
 5. The cell_type should only be 'markdown' if there is no code to add. There may be times where the code has comments. These are still code cells and should have the cell_type 'code'. Any cells that are labeled 'markdown' will be converted to markdown cells by the user.
-6. The assumptions is an optional list of critical assumptions that you made about the data or analysis approach. The assumptions you list here will be displayed to the user so that they can confirm or correct the assumptions. For example: ["NaN values in the impressions column represent 0 impressions", "Only crashes with pedestrian or cyclist fatalities are considered fatal crashes", "Intervention priority combines both volume and severity to identify maximum impact opportunities"].
+6. The analysis_assumptions is an optional list of critical assumptions that you made about the data or analysis approach. The assumptions you list here will be displayed to the user so that they can confirm or correct the assumptions. For example: ["NaN values in the impressions column represent 0 impressions", "Only crashes with pedestrian or cyclist fatalities are considered fatal crashes", "Intervention priority combines both volume and severity to identify maximum impact opportunities"].
 7. Only include important data and analytical assumptions that if incorrect would fundamentally change your analysis conclusions. These should be data handling decisions, methodological choices, and definitional boundaries. Do not include: obvious statements ("Each record is counted once"), result interpretation guidance ("Gaps in the plot represent zero values"), display choices ("Data is sorted for clarity"), internal reasoning ("Bar chart is better than line plot"), or environment assumptions ("Library X is installed"). Prioritize quality over quantity - include only the most critical assumptions or omit the field entirely if there are no critical assumptions made in this step that have not already be shared with the user. If you ever doubt whether an assumption is critical enough to be shared with the user as an assumption, don't include it. Most messages should not include an assumption. 
 8. Do not include the same assumption or variations of the same assumption multiple times in the same conversation. Once you have presented the assumption to the user, they will already have the opportunity to confirm or correct it so do not include it again.
 

mito_ai/completions/providers.py

@@ -160,7 +160,7 @@ This attribute is observed by the websocket provider to push the error to the cl
                 # If we hit a free tier limit, then raise an exception right away without retrying.
                 self.log.exception(f"Error during request_completions: {e}")
                 self.last_error = CompletionError.from_exception(e)
-                log_ai_completion_error('user_key' if self.key_type != MITO_SERVER_KEY else 'mito_server_key', message_type, e)
+                log_ai_completion_error('user_key' if self.key_type != MITO_SERVER_KEY else 'mito_server_key', thread_id or "", message_type, e)
                 raise
 
             except BaseException as e:
@@ -169,14 +169,14 @@ This attribute is observed by the websocket provider to push the error to the cl
                     # Exponential backoff: wait 2^attempt seconds
                     wait_time = 2 ** attempt
                     self.log.info(f"Retrying request_completions after {wait_time}s (attempt {attempt + 1}/{max_retries + 1}): {str(e)}")
-                    log_ai_completion_retry('user_key' if self.key_type != MITO_SERVER_KEY else 'mito_server_key', message_type, e)
+                    log_ai_completion_retry('user_key' if self.key_type != MITO_SERVER_KEY else 'mito_server_key', thread_id or "", message_type, e)
                     await asyncio.sleep(wait_time)
                     continue
                 else:
                     # Final failure after all retries - set error state and raise
                     self.log.exception(f"Error during request_completions after {attempt + 1} attempts: {e}")
                     self.last_error = CompletionError.from_exception(e)
-                    log_ai_completion_error('user_key' if self.key_type != MITO_SERVER_KEY else 'mito_server_key', message_type, e)
+                    log_ai_completion_error('user_key' if self.key_type != MITO_SERVER_KEY else 'mito_server_key', thread_id or "", message_type, e)
                     raise
         
         # This should never be reached due to the raise in the except block,
@@ -264,14 +264,8 @@ This attribute is observed by the websocket provider to push the error to the cl
         except BaseException as e:
             self.log.exception(f"Error during stream_completions: {e}")
             self.last_error = CompletionError.from_exception(e)
-            log(
-                MITO_AI_COMPLETION_ERROR,
-                params={
-                    KEY_TYPE_PARAM: self.key_type,
-                    'message_type': message_type.value,
-                },
-                error=e
-            )
+            log_ai_completion_error('user_key' if self.key_type != MITO_SERVER_KEY else 'mito_server_key', thread_id, message_type, e)
+
             # Send error message to client before raising
             reply_fn(CompletionStreamChunk(
                 parent_id=message_id,

mito_ai/streamlit_conversion/streamlit_agent_handler.py

@@ -4,7 +4,7 @@
 import logging
 import os
 from anthropic.types import MessageParam
-from typing import List, Optional, Tuple, cast, Union
+from typing import List, Optional, Tuple, cast
 
 from mito_ai.logger import get_logger
 from mito_ai.streamlit_conversion.agent_utils import apply_patch_to_text, extract_todo_placeholders, fix_diff_headers

mito_ai/streamlit_conversion/streamlit_utils.py

@@ -60,7 +60,17 @@ def create_app_file(app_directory: str, code: str) -> Tuple[bool, str, str]:
         return False, '', f"Error creating file: {str(e)}"
     except Exception as e:
         return False, '', f"Unexpected error: {str(e)}"
+    
 
+def get_app_path(app_directory: str) -> Optional[str]:
+    """
+    Check if the app.py file exists in the given directory.
+    """
+    app_path = os.path.join(app_directory, "app.py")
+    if not os.path.exists(app_path):
+        return None
+    return app_path
+    
 
 def parse_jupyter_notebook_to_extract_required_content(notebook_path: str) -> Dict[str, Any]:
     """