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

mito_ai/tests/streamlit_conversion/test_streamlit_utils.py

@@ -11,6 +11,7 @@ from mito_ai.streamlit_conversion.streamlit_utils import (
     create_app_file,
     parse_jupyter_notebook_to_extract_required_content
 )
+from mito_ai.path_utils import AbsoluteAppPath, AbsoluteNotebookDirPath, AbsoluteNotebookPath, get_absolute_notebook_path
 from typing import Dict, Any
 
 class TestExtractCodeBlocks:
@@ -48,55 +49,47 @@ class TestCreateAppFile:
 
     def test_create_app_file_success(self, tmp_path):
         """Test successful file creation"""
-        file_path = str(tmp_path)
+        app_path = os.path.join(str(tmp_path), "app.py")
         code = "import streamlit\nst.title('Test')"
         
-        success, app_path, message = create_app_file(file_path, code)
+        create_app_file(AbsoluteAppPath(app_path), code)
         
-        assert success is True
-        assert "Successfully created" in message
+        assert app_path is not None
+        assert os.path.exists(app_path)
         
         # Verify file was created with correct content
-        app_file_path = os.path.join(file_path, "app.py")
-        assert os.path.exists(app_file_path)
-        
-        with open(app_file_path, 'r') as f:
+        with open(app_path, 'r') as f:
             content = f.read()
         assert content == code
 
     def test_create_app_file_io_error(self):
         """Test file creation with IO error"""
-        file_path = "/nonexistent/path/that/should/fail"
+        file_path = AbsoluteAppPath("/nonexistent/path/that/should/fail")
         code = "import streamlit"
         
-        success, app_path, message = create_app_file(file_path, code)
-        
-        assert success is False
-        assert "Error creating file" in message
+        with pytest.raises(Exception):
+            create_app_file(file_path, code)
 
     @patch('builtins.open', side_effect=Exception("Unexpected error"))
     def test_create_app_file_unexpected_error(self, mock_open):
         """Test file creation with unexpected error"""
-        file_path = "/tmp/test"
+        app_path = AbsoluteAppPath("/tmp/test")
         code = "import streamlit"
         
-        success, app_path, message = create_app_file(file_path, code)
-        
-        assert success is False
-        assert "Unexpected error" in message
+        with pytest.raises(Exception, match="Unexpected error"):
+            create_app_file(app_path, code)
 
     def test_create_app_file_empty_code(self, tmp_path):
         """Test creating file with empty code"""
-        file_path = str(tmp_path)
+        app_path = AbsoluteAppPath(os.path.join(str(tmp_path), "app.py"))
         code = ""
         
-        success, app_path, message = create_app_file(file_path, code)
+        create_app_file(app_path, code)
         
-        assert success is True
-        assert "Successfully created" in message
+        assert app_path is not None
+        assert os.path.exists(app_path)
         
-        app_file_path = os.path.join(file_path, "app.py")
-        with open(app_file_path, 'r') as f:
+        with open(app_path, 'r') as f:
             content = f.read()
         assert content == ""
 
@@ -128,7 +121,8 @@ class TestParseJupyterNotebookToExtractRequiredContent:
         with open(notebook_path, 'w') as f:
             json.dump(notebook_data, f)
         
-        result = parse_jupyter_notebook_to_extract_required_content(str(notebook_path))
+        absolute_path = get_absolute_notebook_path(str(notebook_path))
+        result = parse_jupyter_notebook_to_extract_required_content(absolute_path)
         
         # Check that only cell_type and source are preserved
         assert len(result) == 2
@@ -143,8 +137,9 @@ class TestParseJupyterNotebookToExtractRequiredContent:
 
     def test_parse_notebook_file_not_found(self):
         """Test parsing non-existent notebook file"""
-        with pytest.raises(FileNotFoundError, match="Notebook file not found"):
-            parse_jupyter_notebook_to_extract_required_content("/nonexistent/path/notebook.ipynb")
+        from mito_ai.utils.error_classes import StreamlitConversionError
+        with pytest.raises(StreamlitConversionError, match="Notebook file not found"):
+            parse_jupyter_notebook_to_extract_required_content(AbsoluteNotebookPath("/nonexistent/path/notebook.ipynb"))
 
     def test_parse_notebook_with_missing_cell_fields(self, tmp_path):
         """Test parsing notebook where cells are missing cell_type or source"""
@@ -169,7 +164,8 @@ class TestParseJupyterNotebookToExtractRequiredContent:
         with open(notebook_path, 'w') as f:
             json.dump(notebook_data, f)
         
-        result = parse_jupyter_notebook_to_extract_required_content(str(notebook_path))
+        absolute_path = get_absolute_notebook_path(str(notebook_path))
+        result = parse_jupyter_notebook_to_extract_required_content(absolute_path)
         
         assert len(result) == 3
         assert result[0]['cell_type'] == 'code'
@@ -191,6 +187,7 @@ class TestParseJupyterNotebookToExtractRequiredContent:
         with open(notebook_path, 'w') as f:
             json.dump(notebook_data, f)
         
-        result = parse_jupyter_notebook_to_extract_required_content(str(notebook_path))
+        absolute_path = get_absolute_notebook_path(str(notebook_path))
+        result = parse_jupyter_notebook_to_extract_required_content(absolute_path)
         
         assert result == []

mito_ai/tests/streamlit_conversion/test_validate_streamlit_app.py

@@ -7,10 +7,10 @@ from unittest.mock import patch, MagicMock
 from mito_ai.streamlit_conversion.validate_streamlit_app import (
     get_syntax_error,
     get_runtime_errors,
-    check_for_errors,
     validate_app
 )
 import pytest
+from mito_ai.path_utils import AbsoluteNotebookPath
 
 
 class TestGetSyntaxError:
@@ -19,33 +19,33 @@ class TestGetSyntaxError:
     @pytest.mark.parametrize("code,expected_error,test_description", [
         # Valid Python code should return no error
         (
-            "import streamlit\nst.title('Hello World')",
-            None,
-            "valid Python code"
+                "import streamlit\nst.title('Hello World')",
+                None,
+                "valid Python code"
         ),
         # Invalid Python syntax should be caught
         (
-            "import streamlit\nst.title('Hello World'",
-            "SyntaxError",
-            "invalid Python code"
+                "import streamlit\nst.title('Hello World'",
+                "SyntaxError",
+                "invalid Python code"
         ),
         # Empty streamlit app is valid
         (
-            "",
-            None,
-            "empty code"
+                "",
+                None,
+                "empty code"
         ),
     ])
     def test_get_syntax_error(self, code, expected_error, test_description):
         """Test syntax validation with various code inputs"""
         error = get_syntax_error(code)
-        
+
         if expected_error is None:
             assert error is None, f"Expected no error for {test_description}"
         else:
             assert error is not None, f"Expected error for {test_description}"
             assert expected_error in error, f"Expected '{expected_error}' in error for {test_description}"
-        
+
 class TestGetRuntimeErrors:
     """Test cases for get_runtime_errors function"""
 
@@ -56,7 +56,9 @@ class TestGetRuntimeErrors:
     ])
     def test_get_runtime_errors(self, app_code, expected_error):
         """Test getting runtime errors"""
-        errors = get_runtime_errors(app_code, '/app.py')
+        
+        absolute_path = AbsoluteNotebookPath('/notebook.ipynb')
+        errors = get_runtime_errors(app_code, absolute_path)
         
         if expected_error is None:
             assert errors is None
@@ -85,23 +87,26 @@ df=pd.read_csv('data.csv')
             with open(csv_path, "w") as f:
                 f.write("name,age\nJohn,25\nJane,30")
                
-            errors = get_runtime_errors(app_code, app_path)
+            errors = get_runtime_errors(app_code, AbsoluteNotebookPath(app_path))
             assert errors is None
 
 class TestValidateApp:
     """Test cases for validate_app function"""
 
-    @pytest.mark.parametrize("app_code,expected_has_validation_error,expected_error_message", [
+    @pytest.mark.parametrize("app_code,expected_has_errors,expected_error_message", [
         ("x=5", False, ""),
         ("1/0", True, "division by zero"),
-        ("print('Hello World'", True, "SyntaxError"),
+        # Syntax errors are caught during runtime
         ("", False, ""),
     ])
-    def test_validate_app(self, app_code, expected_has_validation_error, expected_error_message):
+    def test_validate_app(self, app_code, expected_has_errors, expected_error_message):
         """Test the validate_app function"""
-        has_validation_error, errors = validate_app(app_code, '/app.py')
+        errors = validate_app(app_code, AbsoluteNotebookPath('/notebook.ipynb'))
         
-        assert has_validation_error == expected_has_validation_error
-        assert expected_error_message in str(errors)
+        has_errors = len(errors) > 0
+        assert has_errors == expected_has_errors
+        if expected_error_message:
+            errors_str = str(errors)
+            assert expected_error_message in errors_str
         
     

mito_ai/tests/streamlit_preview/test_streamlit_preview_handler.py

@@ -4,85 +4,110 @@
 import pytest
 import os
 import tempfile
-from unittest.mock import patch, AsyncMock, MagicMock
-from mito_ai.streamlit_preview.utils import ensure_app_exists
+from unittest.mock import patch, Mock, AsyncMock, MagicMock
+from mito_ai.streamlit_preview.handlers import StreamlitPreviewHandler
+from mito_ai.path_utils import AbsoluteNotebookPath
 
 
-class TestEnsureAppExists:
-    """Test cases for the ensure_app_exists function."""
+class TestStreamlitPreviewHandler:
+    """Test cases for the StreamlitPreviewHandler."""
 
     @pytest.mark.asyncio
     @pytest.mark.parametrize(
-        "app_exists,streamlit_handler_success,expected_success,expected_error,streamlit_handler_called,streamlit_handler_return",
+        "app_exists,force_recreate,streamlit_handler_called",
         [
-            # 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)
-            ),
-            # Test case 2: App doesn't exist, streamlit_handler succeeds
-            (
-                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
-            )
+            # Test case 1: App exists, not forcing recreate - should not call streamlit_handler
+            (True, False, False),
+            # Test case 2: App doesn't exist - should call streamlit_handler
+            (False, False, True),
+            # Test case 3: App exists but forcing recreate - should call streamlit_handler
+            (True, True, True),
         ],
         ids=[
-            "app_exists_uses_existing_file",
-            "app_does_not_exist_generates_new_one_success",
+            "app_exists_no_force_recreate",
+            "app_does_not_exist_generates_new_one",
+            "app_exists_force_recreate",
         ]
     )
-    async def test_ensure_app_exists(
+    async def test_post_handler_app_generation(
         self,
         app_exists,
-        streamlit_handler_success,
-        expected_success,
-        expected_error,
+        force_recreate,
         streamlit_handler_called,
-        streamlit_handler_return,
     ):
-        """Test ensure_app_exists function with various scenarios."""
+        """Test StreamlitPreviewHandler POST method with various scenarios."""
         with tempfile.TemporaryDirectory() as temp_dir:
             notebook_path = os.path.join(temp_dir, "test_notebook.ipynb")
+            app_path = os.path.join(temp_dir, "app.py")
             
-            # Set up app_path based on whether app exists
-            app_path = os.path.join(temp_dir, "app.py") if app_exists else None
+            # Create notebook file
+            with open(notebook_path, "w") as f:
+                f.write('{"cells": []}')
             
             # Create app.py file if it should exist
             if app_exists:
-                assert app_path is not None  # Type assertion for mypy
                 with open(app_path, "w") as f:
                     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:
-                mock_get_app_path.return_value = app_path
+            # Create a properly mocked Tornado application with required attributes
+            mock_application = MagicMock()
+            mock_application.ui_methods = {}
+            mock_application.ui_modules = {}
+            mock_application.settings = {}
+            
+            # Create a mock request with necessary tornado setup
+            mock_request = MagicMock()
+            mock_request.connection = MagicMock()
+            mock_request.connection.context = MagicMock()
+            
+            # Create handler instance  
+            handler = StreamlitPreviewHandler(
+                application=mock_application,
+                request=mock_request,
+            )
+            handler.initialize()
+            
+            # Mock authentication - set current_user to bypass @tornado.web.authenticated
+            handler.current_user = "test_user"  # type: ignore
+            
+            # Mock the finish method and other handler methods
+            finish_called = []
+            def mock_finish_func(response):
+                finish_called.append(response)
+            
+            # Mock streamlit_handler and preview manager
+            with patch.object(handler, 'get_json_body', return_value={
+                "notebook_path": notebook_path,
+                "force_recreate": force_recreate,
+                "edit_prompt": ""
+            }), \
+                 patch.object(handler, 'finish', side_effect=mock_finish_func), \
+                 patch.object(handler, 'set_status'), \
+                 patch('mito_ai.streamlit_preview.handlers.streamlit_handler', new_callable=AsyncMock) as mock_streamlit_handler, \
+                 patch.object(handler.preview_manager, 'start_streamlit_preview', return_value=8501) as mock_start_preview, \
+                 patch('mito_ai.streamlit_preview.handlers.log_streamlit_app_preview_success'):
+                
+                # Call the handler
+                await handler.post()  # type: ignore[misc]
+                
+                # Verify streamlit_handler was called or not called as expected
+                if streamlit_handler_called:
+                    assert mock_streamlit_handler.called
+                    # Verify it was called with the absolute notebook path
+                    call_args = mock_streamlit_handler.call_args
+                    assert call_args[0][0] == notebook_path  # First argument should be the notebook path
+                    assert call_args[0][1] == ""  # Second argument should be the edit_prompt
+                else:
+                    mock_streamlit_handler.assert_not_called()
+                
+                # Verify preview was started
+                mock_start_preview.assert_called_once()
                 
-                # 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
-                    
-                    # Verify get_app_path was called with the correct directory
-                    mock_get_app_path.assert_called_once_with(temp_dir)
-                    
-                    # Verify streamlit_handler was called or not called as expected
-                    if streamlit_handler_called:
-                        mock_streamlit_handler.assert_called_once_with(notebook_path, "")
-                    else:
-                        mock_streamlit_handler.assert_not_called()
+                # Verify response was sent
+                assert len(finish_called) == 1
+                response = finish_called[0]
+                assert response["type"] == "success"
+                assert "port" in response
+                assert "id" in response
 
     

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, 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_preview")
+            
+            assert expected_message in str(exc_info.value).lower()
     
     @pytest.mark.parametrize("preview_id,expected_result", [
         ("existing_preview", True),
@@ -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", 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: