mito-ai 0.1.46__py3-none-any.whl → 0.1.49__py3-none-any.whl

mito_ai/tests/streamlit_preview/test_streamlit_preview_manager.py

@@ -15,11 +15,8 @@ from typing import Any
 
 from mito_ai.streamlit_preview.manager import (
     StreamlitPreviewManager, 
-    PreviewProcess, 
-    get_preview_manager
+    PreviewProcess
 )
-from mito_ai.streamlit_preview.handlers import StreamlitPreviewHandler
-from mito_ai.streamlit_conversion.streamlit_utils import get_app_path
 
 
 class TestStreamlitPreviewManager:
@@ -81,25 +78,22 @@ st.write("Hello, World!")
             mock_requests_get.return_value = mock_response
             
             # Test
-            success, message, port = manager.start_streamlit_preview(app_directory, preview_id)
+            port = manager.start_streamlit_preview(app_directory, 'test-file-name.py',  preview_id)
             
             # Assertions
-            assert success == expected_success
-            if expected_success:
-                assert "successfully" in message.lower()
-                assert isinstance(port, int)
-                assert port > 0
-                
-                # Verify subprocess was called correctly
-                mock_popen.assert_called_once()
-                call_args = mock_popen.call_args
-                assert "streamlit" in call_args[0][0]
-                assert "run" in call_args[0][0]
-                assert "--server.headless" in call_args[0][0]
-                assert "--server.address" in call_args[0][0]
-                
-                # Cleanup
-                manager.stop_preview(preview_id)
+            assert isinstance(port, int)
+            assert port > 0
+            
+            # Verify subprocess was called correctly
+            mock_popen.assert_called_once()
+            call_args = mock_popen.call_args
+            assert "streamlit" in call_args[0][0]
+            assert "run" in call_args[0][0]
+            assert "--server.headless" in call_args[0][0]
+            assert "--server.address" in call_args[0][0]
+            
+            # Cleanup
+            manager.stop_preview(preview_id)
     
     @pytest.mark.parametrize("exception_type,expected_message", [
         (Exception("Temp dir creation failed"), "failed to start preview"),
@@ -108,13 +102,15 @@ st.write("Hello, World!")
     ])
     def test_start_streamlit_preview_exceptions(self, manager, sample_app_code, exception_type, expected_message):
         """Test streamlit preview start with different exceptions."""
-        with patch('tempfile.mkdtemp', side_effect=exception_type):
+        from mito_ai.utils.error_classes import StreamlitPreviewError
+        
+        with patch('subprocess.Popen', side_effect=exception_type):
             app_directory = "/tmp/test_dir"
-            success, message, port = manager.start_streamlit_preview(app_directory, "test_preview")
             
-            assert success is False
-            assert expected_message in message.lower()
-            assert port is None
+            with pytest.raises(StreamlitPreviewError) as exc_info:
+                manager.start_streamlit_preview(app_directory, 'test-file-name.py',  "test_preview")
+            
+            assert expected_message in str(exc_info.value).lower()
     
     @pytest.mark.parametrize("preview_id,expected_result", [
         ("existing_preview", True),
@@ -146,7 +142,7 @@ st.write("Hello, World!")
                 mock_open.return_value.__enter__.return_value = mock_file
                 mock_exists.return_value = True
                 
-                manager.start_streamlit_preview(app_directory, preview_id)
+                manager.start_streamlit_preview(app_directory, 'test-file-name.py',  preview_id)
     
     @pytest.mark.parametrize("process_behavior,expected_kill_called", [
         (subprocess.TimeoutExpired("cmd", 5), True),
@@ -179,7 +175,7 @@ st.write("Hello, World!")
             mock_exists.return_value = True
             
             # Start a preview
-            manager.start_streamlit_preview(app_directory, "test_preview")
+            manager.start_streamlit_preview(app_directory, 'test-file-name.py',  "test_preview")
             
             # Setup process behavior for stop
             if process_behavior:
@@ -219,7 +215,7 @@ st.write("Hello, World!")
                 mock_open.return_value.__enter__.return_value = mock_file
                 mock_exists.return_value = True
                 
-                manager.start_streamlit_preview("/tmp/test_dir", preview_id)
+                manager.start_streamlit_preview("/tmp/test_dir", 'test-file-name.py', preview_id)
         
         preview = manager.get_preview(preview_id)
         
@@ -245,15 +241,7 @@ st.write("Hello, World!")
         
         assert preview.proc == proc
         assert preview.port == port
-    
-    def test_get_preview_manager_singleton(self):
-        """Test that get_preview_manager returns the same instance."""
-        manager1 = get_preview_manager()
-        manager2 = get_preview_manager()
-        
-        assert manager1 is manager2
-        assert isinstance(manager1, StreamlitPreviewManager)
-    
+
     @pytest.mark.parametrize("num_previews", [1, 2, 3])
     def test_concurrent_previews(self, manager, sample_app_code, num_previews):
         """Test managing multiple concurrent previews."""
@@ -284,8 +272,7 @@ st.write("Hello, World!")
             
             # Start multiple previews
             for preview_id in preview_ids:
-                success, _, port = manager.start_streamlit_preview("/tmp/test_dir", preview_id)
-                assert success is True
+                port = manager.start_streamlit_preview("/tmp/test_dir", 'test-file-name.py', preview_id)
                 ports.append(port)
             
             # Assertions

mito_ai/user/handlers.py

@@ -3,9 +3,11 @@
 
 import json
 import tornado
+from typing import Any, Optional
 from jupyter_server.base.handlers import APIHandler
 from mito_ai.utils.db import get_user_field, set_user_field
 from mito_ai.utils.telemetry_utils import identify
+from mito_ai.utils.version_utils import is_pro
 
 
 class UserHandler(APIHandler):
@@ -13,7 +15,15 @@ class UserHandler(APIHandler):
 
     @tornado.web.authenticated
     def get(self, key: str) -> None:
-        value = get_user_field(key)
+        value: Optional[Any] = None
+
+        if key == "is_pro":
+            # Special case, since we don't store this key 
+            # in the user.json file.
+            value = str(is_pro())
+        else: 
+            value = get_user_field(key)
+
         if value is None:
             self.set_status(404)
             self.finish(json.dumps({"error": f"User field with key '{key}' not found"}))
@@ -29,5 +39,7 @@ class UserHandler(APIHandler):
             return
 
         set_user_field(key, data["value"])
-        identify() # Log the new user 
-        self.finish(json.dumps({"status": "success", "key": key, "value": data["value"]}))
+        identify()  # Log the new user
+        self.finish(
+            json.dumps({"status": "success", "key": key, "value": data["value"]})
+        )

mito_ai/utils/create.py

@@ -36,6 +36,17 @@ def is_user_json_exists_and_valid_json() -> bool:
     except:
         return False
 
+def get_temp_user_id() -> Optional[str]:
+    """
+    Looks for a temporary user ID, generated by the desktop app.
+    """
+    temp_user_id_path = os.path.join(MITO_FOLDER, 'temp_user_id.txt')
+
+    if os.path.exists(temp_user_id_path):
+        with open(temp_user_id_path, 'r') as f:
+            return f.read()
+
+    return None
 
 def try_create_user_json_file() -> None:
     
@@ -50,7 +61,12 @@ def try_create_user_json_file() -> None:
         with open(USER_JSON_PATH, 'w+') as f:
             f.write(json.dumps(USER_JSON_DEFAULT))
 
-        # Then, we take special care to put all the testing/CI environments 
+        # Next, look for a temp user id
+        temp_user_id = get_temp_user_id()
+        if temp_user_id:
+            set_user_field(UJ_STATIC_USER_ID, temp_user_id)
+
+        # Finally, we take special care to put all the testing/CI environments 
         # (e.g. Github actions) under one ID and email
         if is_running_test():
             set_user_field(UJ_STATIC_USER_ID, GITHUB_ACTION_ID)

mito_ai/utils/error_classes.py

@@ -0,0 +1,42 @@
+# Copyright (c) Saga Inc.
+# Distributed under the terms of the GNU Affero General Public License v3.0 License.
+
+from mito_ai.app_deploy.models import AppDeployError
+
+class MitoAppError(Exception):
+    """Exception raised for custom error in the application."""
+
+    def __init__(self, message: str, error_code: int) -> None:
+        super().__init__(message)
+        self.message = message
+        self.error_code = error_code
+
+class StreamlitPreviewError(MitoAppError):
+    def __str__(self) -> str:
+        return f"[PreviewError]: {self.message} (Error Code: {self.error_code})"
+
+class StreamlitConversionError(MitoAppError):
+    def __str__(self) -> str:
+        return f"[ConversionError]: {self.message} (Error Code: {self.error_code})"
+
+
+class StreamlitDeploymentError(MitoAppError):
+    """Raised when a deployment operation fails."""
+
+    def __init__(self, error: AppDeployError):
+        self.error = error
+        self.error_type = error.error_type
+        self.message_id = getattr(error, "message_id", "ErrorMessageID")
+        self.error_code = getattr(error, "error_code", 500)
+        self.hint = getattr(error, "hint", "")
+        self.traceback = getattr(error, "traceback", "")
+        self.error_type = getattr(error, "error_type", "Error")
+        self.message = error.message
+        print(f"self_message: {self.message}")
+        super().__init__(self.message, self.error_code)
+
+    def __str__(self) -> str:
+        base = f"[DeploymentError]: {self.message} (Error Code: {self.error_code})"
+        if self.hint:
+            base += f"\nHint: {self.hint}"
+        return base

mito_ai/utils/message_history_utils.py

@@ -10,6 +10,7 @@ from mito_ai.completions.prompt_builders.prompt_constants import (
     ACTIVE_CELL_OUTPUT_SECTION_HEADING,
     GET_CELL_OUTPUT_TOOL_RESPONSE_SECTION_HEADING,
     FILES_SECTION_HEADING,
+    STREAMLIT_APP_STATUS_SECTION_HEADING,
     VARIABLES_SECTION_HEADING,
     JUPYTER_NOTEBOOK_SECTION_HEADING,
     CONTENT_REMOVED_PLACEHOLDER
@@ -31,7 +32,8 @@ def trim_sections_from_message_content(content: str) -> str:
         JUPYTER_NOTEBOOK_SECTION_HEADING,
         GET_CELL_OUTPUT_TOOL_RESPONSE_SECTION_HEADING,
         ACTIVE_CELL_OUTPUT_SECTION_HEADING,
-        ACTIVE_CELL_ID_SECTION_HEADING
+        ACTIVE_CELL_ID_SECTION_HEADING,
+        STREAMLIT_APP_STATUS_SECTION_HEADING
     ]
     
     for heading in section_headings:

mito_ai/utils/telemetry_utils.py

@@ -3,7 +3,7 @@
 
 import json
 import os
-from typing import Any, Dict, Literal, Optional
+from typing import Any, Dict, Literal, Optional, List
 from mito_ai.utils.version_utils import MITOSHEET_HELPER_PRIVATE, is_pro
 from mito_ai.utils.schema import UJ_AI_MITO_API_NUM_USAGES, UJ_MITOSHEET_TELEMETRY, UJ_STATIC_USER_ID, UJ_USER_EMAIL, UJ_FEEDBACKS_V2
 from mito_ai.utils.db import get_user_field
@@ -190,7 +190,24 @@ def log(
 
     if thread_id is not None:
         final_params['thread_id'] = thread_id
-
+        
+    # Process parameters that need chunking
+    params_to_remove = []
+    params_to_add = {}
+    
+    for param_name, param_value in final_params.items():
+        if isinstance(param_value, str) and len(param_value) > 250:
+            # Mark for removal
+            params_to_remove.append(param_name)
+            # Get chunked parameters
+            chunked_params = chunk_param(param_value, param_name)
+            params_to_add.update(chunked_params)
+    
+    # Apply the changes
+    for param_name in params_to_remove:
+        del final_params[param_name]
+    final_params.update(params_to_add)
+            
     # Finally, do the acutal logging. We do not log anything when tests are
     # running, or if telemetry is turned off
     if not is_running_test() and telemetry_turned_on(key_type):
@@ -368,7 +385,7 @@ def log_file_upload_failure(error: str) -> None:
     log("mito_ai_file_upload_failure", params={"error_message": error})
 
 def log_ai_completion_retry(key_type: Literal['mito_server_key', 'user_key'], thread_id: str, message_type: MessageType, error: BaseException) -> None:
-    log(MITO_AI_COMPLETION_RETRY, params={KEY_TYPE_PARAM: key_type, "message_type": message_type}, thread_id=thread_id, key_type=key_type, error=error)
+    log(MITO_AI_COMPLETION_RETRY, params={KEY_TYPE_PARAM: key_type, "message_type": message_type.value}, thread_id=thread_id, key_type=key_type, error=error)
     
 def log_ai_completion_error(
     key_type: Literal['mito_server_key', 'user_key'],
@@ -376,40 +393,88 @@ def log_ai_completion_error(
     message_type: MessageType, 
     error: BaseException
 ) -> None:
-    log(MITO_AI_COMPLETION_ERROR, params={KEY_TYPE_PARAM: key_type, "message_type": message_type}, thread_id=thread_id, key_type=key_type, error=error)
+    log(MITO_AI_COMPLETION_ERROR, params={KEY_TYPE_PARAM: key_type, "message_type": message_type.value}, thread_id=thread_id, key_type=key_type, error=error)
     
 def log_mito_server_free_tier_limit_reached(key_type: Literal['mito_server_key', 'user_key'], message_type: MessageType) -> None:
-    log(MITO_SERVER_FREE_TIER_LIMIT_REACHED, params={KEY_TYPE_PARAM: key_type, "message_type": message_type}, key_type=key_type)
+    log(MITO_SERVER_FREE_TIER_LIMIT_REACHED, params={KEY_TYPE_PARAM: key_type, "message_type": message_type.value}, key_type=key_type)
 
 
 #################################
-# Streamlit Conversion
+# Streamlit Logs
 #################################
 
-def log_streamlit_app_creation_success(key_type: Literal['mito_server_key', 'user_key'], message_type: MessageType, edit_prompt: str) -> None:
+### 
+# Converting Notebook into Streamlit App code
+###
+
+def log_streamlit_app_conversion_success(key_type: Literal['mito_server_key', 'user_key'], message_type: MessageType, edit_prompt: str) -> None:
     log(
-        "mito_ai_streamlit_app_creation_success", 
+        "mito_ai_streamlit_app_conversion_success", 
         key_type=key_type,
         params={
             "edit_prompt": edit_prompt,
+            "message_type": message_type.value
         }
     )
     
-def log_streamlit_app_creation_retry(key_type: Literal['mito_server_key', 'user_key'], message_type: MessageType, error: str) -> None:
+def log_streamlit_app_validation_retry(key_type: Literal['mito_server_key', 'user_key'], message_type: MessageType, error: List[str]) -> None:
     log(
-        "mito_ai_streamlit_app_creation_retry", 
+        "mito_ai_streamlit_app_conversion_retry",
         params={
             "error_message": error,
+            "message_type": message_type.value
         },
         key_type=key_type
     )
     
-def log_streamlit_app_creation_error(key_type: Literal['mito_server_key', 'user_key'], message_type: MessageType, error: str, edit_prompt: str) -> None:
+def log_streamlit_app_conversion_error(key_type: Literal['mito_server_key', 'user_key'], message_type: MessageType, error_message: str, formatted_traceback: str, edit_prompt: str) -> None:
     log(
-        "mito_ai_streamlit_app_creation_error", 
+        "mito_ai_streamlit_app_conversion_error",
         params={
-            "error_message": error,
+            "error_message": error_message,
+            "traceback": formatted_traceback,
             "edit_prompt": edit_prompt,
+            "message_type": message_type.value
         },
         key_type=key_type
     )
+    
+###
+# Setting up Preview
+###
+
+def log_streamlit_app_preview_success(key_type: Literal['mito_server_key', 'user_key'], message_type: MessageType, edit_prompt: str) -> None:
+    log(
+        "mito_ai_streamlit_app_preview_success",
+        key_type=key_type,
+        params={
+            "edit_prompt": edit_prompt,
+            "message_type": message_type.value
+        }
+    )
+    
+def log_streamlit_app_preview_failure(key_type: Literal['mito_server_key', 'user_key'], message_type: MessageType, error_message: str, formatted_traceback: str, edit_prompt: str) -> None:
+    log(
+        "mito_ai_streamlit_app_preview_failure",
+        key_type=key_type,
+        params={
+            "error_message": error_message,
+            "traceback": formatted_traceback,
+            "message_type": message_type.value,
+            "edit_prompt": edit_prompt
+        }
+    )
+    
+###
+# Deploying Streamlit App
+###
+
+def log_streamlit_app_deployment_failure(key_type: Literal['mito_server_key', 'user_key'], message_type: MessageType, error: Dict) -> None:
+    log(
+        "mito_ai_streamlit_app_deployment_failure",
+        key_type=key_type,
+        params={
+            "error": error, # Contains all details in app_deploy.models.AppDeployError class
+            "message_type": message_type.value
+        }
+    )