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

mito_ai/tests/streamlit_preview/test_streamlit_preview_handler.py

@@ -4,8 +4,9 @@
 import pytest
 import os
 import tempfile
-from unittest.mock import patch, AsyncMock, MagicMock
+from unittest.mock import patch
 from mito_ai.streamlit_preview.utils import ensure_app_exists
+from mito_ai.path_utils import AbsoluteNotebookPath, get_absolute_notebook_path
 
 
 class TestEnsureAppExists:
@@ -13,14 +14,12 @@ class TestEnsureAppExists:
 
     @pytest.mark.asyncio
     @pytest.mark.parametrize(
-        "app_exists,streamlit_handler_success,expected_success,expected_error,streamlit_handler_called,streamlit_handler_return",
+        "app_exists,streamlit_handler_success,streamlit_handler_called,streamlit_handler_return",
         [
             # Test case 1: App exists, should use existing file
             (
                 True,  # app_exists
                 True,  # streamlit_handler_success (not relevant)
-                True,  # expected_success
-                "",    # expected_error
                 False, # streamlit_handler_called
                 None,  # streamlit_handler_return (not used)
             ),
@@ -28,10 +27,8 @@ class TestEnsureAppExists:
             (
                 False,  # app_exists
                 True,   # streamlit_handler_success
-                True,   # expected_success
-                "",     # expected_error
                 True,   # streamlit_handler_called
-                (True, "/path/to/app.py", "Success"),  # streamlit_handler_return
+                "/path/to/app.py",  # streamlit_handler_return
             )
         ],
         ids=[
@@ -43,8 +40,6 @@ class TestEnsureAppExists:
         self,
         app_exists,
         streamlit_handler_success,
-        expected_success,
-        expected_error,
         streamlit_handler_called,
         streamlit_handler_return,
     ):
@@ -62,19 +57,21 @@ class TestEnsureAppExists:
                     f.write("import streamlit as st\nst.write('Hello World')")
             
             # Mock get_app_path to return the appropriate value
-            with patch('mito_ai.streamlit_preview.utils.get_app_path') as mock_get_app_path:
+            with patch('mito_ai.streamlit_preview.utils.get_absolute_notebook_dir_path') as mock_get_dir_path, \
+                patch('mito_ai.streamlit_preview.utils.get_absolute_app_path') as mock_get_app_path, \
+                patch('mito_ai.streamlit_preview.utils.does_app_path_exists') as mock_app_exists:
+            
+                # Set up mocks
+                mock_get_dir_path.return_value = temp_dir
                 mock_get_app_path.return_value = app_path
+                mock_app_exists.return_value = app_exists
                 
                 # Mock streamlit_handler
                 with patch('mito_ai.streamlit_preview.utils.streamlit_handler') as mock_streamlit_handler:
                     if streamlit_handler_return is not None:
                         mock_streamlit_handler.return_value = streamlit_handler_return
                     
-                    success, error_msg = await ensure_app_exists(notebook_path, False, "")
-                    
-                    # Assertions
-                    assert success == expected_success
-                    assert error_msg == expected_error
+                    await ensure_app_exists(AbsoluteNotebookPath(notebook_path), False, "")
                     
                     # Verify get_app_path was called with the correct directory
                     mock_get_app_path.assert_called_once_with(temp_dir)

mito_ai/tests/streamlit_preview/test_streamlit_preview_manager.py

@@ -18,8 +18,6 @@ from mito_ai.streamlit_preview.manager import (
     PreviewProcess, 
     get_preview_manager
 )
-from mito_ai.streamlit_preview.handlers import StreamlitPreviewHandler
-from mito_ai.streamlit_conversion.streamlit_utils import get_app_path
 
 
 class TestStreamlitPreviewManager:
@@ -81,25 +79,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, 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 +103,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_preview")
+            
+            assert expected_message in str(exc_info.value).lower()
     
     @pytest.mark.parametrize("preview_id,expected_result", [
         ("existing_preview", True),
@@ -284,8 +281,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", 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):
@@ -383,33 +400,81 @@ def log_mito_server_free_tier_limit_reached(key_type: Literal['mito_server_key',
 
 
 #################################
-# 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
         }
     )
     
-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
         },
         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
         },
         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
+        }
+    )
+    
+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, 
+            "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
+        }
+    )

{mito_ai-0.1.46.data → mito_ai-0.1.47.data}/data/share/jupyter/labextensions/mito_ai/build_log.json

@@ -712,7 +712,7 @@
             "semver": {},
             "vscode-diff": {},
             "mito_ai": {
-              "version": "0.1.46",
+              "version": "0.1.47",
               "singleton": true,
               "import": "/home/runner/work/mito/mito/mito-ai/lib/index.js"
             }

{mito_ai-0.1.46.data → mito_ai-0.1.47.data}/data/share/jupyter/labextensions/mito_ai/package.json

@@ -1,6 +1,6 @@
 {
   "name": "mito_ai",
-  "version": "0.1.46",
+  "version": "0.1.47",
   "description": "AI chat for JupyterLab",
   "keywords": [
     "jupyter",
@@ -140,7 +140,7 @@
     "outputDir": "mito_ai/labextension",
     "schemaDir": "schema",
     "_build": {
-      "load": "static/remoteEntry.54126ab6511271265443.js",
+      "load": "static/remoteEntry.e22c6cd4e56c32116daa.js",
       "extension": "./extension",
       "style": "./style"
     }

{mito_ai-0.1.46.data → mito_ai-0.1.47.data}/data/share/jupyter/labextensions/mito_ai/schemas/mito_ai/package.json.orig

@@ -1,6 +1,6 @@
 {
   "name": "mito_ai",
-  "version": "0.1.46",
+  "version": "0.1.47",
   "description": "AI chat for JupyterLab",
   "keywords": [
     "jupyter",