mito-ai 0.1.37__py3-none-any.whl → 0.1.38__py3-none-any.whl

mito_ai/streamlit_preview/urls.py

@@ -0,0 +1,22 @@
+# Copyright (c) Saga Inc.
+# Distributed under the terms of the GNU Affero General Public License v3.0 License.
+
+from typing import Any, List, Tuple
+from jupyter_server.utils import url_path_join
+from mito_ai.streamlit_preview.handlers import StreamlitPreviewHandler
+
+def get_streamlit_preview_urls(base_url: str) -> List[Tuple[str, Any, dict]]:
+    """Get all streamlit preview related URL patterns.
+
+    Args:
+        base_url: The base URL for the Jupyter server
+
+    Returns:
+        List of (url_pattern, handler_class, handler_kwargs) tuples
+    """
+    BASE_URL = base_url + "/mito-ai"
+    
+    return [
+        (url_path_join(BASE_URL, "streamlit-preview"), StreamlitPreviewHandler, {}),
+        (url_path_join(BASE_URL, "streamlit-preview/(.+)"), StreamlitPreviewHandler, {}),
+    ]

mito_ai/tests/streamlit_conversion/test_streamlit_agent_handler.py

@@ -170,19 +170,19 @@ class TestStreamlitHandler:
         mock_validator.return_value = (False, "")
         
         # Mock file creation
-        mock_create_file.return_value = (True, "File created successfully")
+        mock_create_file.return_value = (True, "/path/to/app", "File created successfully")
         
-        result = await streamlit_handler("/path/to/notebook.ipynb", "/path/to/app")
+        result = await streamlit_handler("/path/to/notebook.ipynb")
         
         assert result[0] is True
-        assert "File created successfully" in result[1]
+        assert "File created successfully" in result[2]
         
         # Verify calls
         mock_parse.assert_called_once_with("/path/to/notebook.ipynb")
         mock_generator_class.assert_called_once_with(mock_notebook_data)
         mock_generator.generate_streamlit_code.assert_called_once()
         mock_validator.assert_called_once_with("import streamlit\nst.title('Test')")
-        mock_create_file.assert_called_once_with("/path/to/app", "import streamlit\nst.title('Test')")
+        mock_create_file.assert_called_once_with("/path/to", "import streamlit\nst.title('Test')")
 
     @pytest.mark.asyncio
     @patch('mito_ai.streamlit_conversion.streamlit_agent_handler.parse_jupyter_notebook_to_extract_required_content')
@@ -193,20 +193,21 @@ class TestStreamlitHandler:
         # Mock notebook parsing
         mock_notebook_data: dict = {"cells": []}
         mock_parse.return_value = mock_notebook_data
-        
+    
         # Mock code generation
         mock_generator = AsyncMock()
         mock_generator.generate_streamlit_code.return_value = "import streamlit\nst.title('Test')"
         mock_generator.correct_error_in_generation.return_value = "import streamlit\nst.title('Fixed')"
         mock_generator_class.return_value = mock_generator
-        
-        # Mock validation (always errors)
+    
+        # Mock validation (always errors) - FIX: Return only 2 values
         mock_validator.return_value = (True, "Persistent error")
+    
+        result = await streamlit_handler("/path/to/notebook.ipynb")
         
-        result = await streamlit_handler("/path/to/notebook.ipynb", "/path/to/app")
-        
+        # Verify the result indicates failure
         assert result[0] is False
-        assert "Error generating streamlit code by agent" in result[1]
+        assert "Error generating streamlit code by agent" in result[2]
         
         # Verify that error correction was called 5 times (max retries)
         assert mock_generator.correct_error_in_generation.call_count == 5
@@ -231,12 +232,12 @@ class TestStreamlitHandler:
         mock_validator.return_value = (False, "")
         
         # Mock file creation failure
-        mock_create_file.return_value = (False, "Permission denied")
+        mock_create_file.return_value = (False, None, "Permission denied")
         
-        result = await streamlit_handler("/path/to/notebook.ipynb", "/path/to/app")
+        result = await streamlit_handler("/path/to/notebook.ipynb")
         
         assert result[0] is False
-        assert "Permission denied" in result[1]
+        assert "Permission denied" in result[2]
 
     @pytest.mark.asyncio
     @patch('mito_ai.streamlit_conversion.streamlit_agent_handler.parse_jupyter_notebook_to_extract_required_content')
@@ -245,7 +246,7 @@ class TestStreamlitHandler:
         mock_parse.side_effect = FileNotFoundError("Notebook not found")
         
         with pytest.raises(FileNotFoundError, match="Notebook not found"):
-            await streamlit_handler("/path/to/notebook.ipynb", "/path/to/app")
+            await streamlit_handler("/path/to/notebook.ipynb")
 
     @pytest.mark.asyncio
     @patch('mito_ai.streamlit_conversion.streamlit_agent_handler.parse_jupyter_notebook_to_extract_required_content')
@@ -262,4 +263,4 @@ class TestStreamlitHandler:
         mock_generator_class.return_value = mock_generator
         
         with pytest.raises(Exception, match="Generation failed"):
-            await streamlit_handler("/path/to/notebook.ipynb", "/path/to/app")
+            await streamlit_handler("/path/to/notebook.ipynb")

mito_ai/tests/streamlit_conversion/test_streamlit_utils.py

@@ -13,7 +13,6 @@ from mito_ai.streamlit_conversion.streamlit_utils import (
 )
 from typing import Dict, Any
 
-
 class TestExtractCodeBlocks:
     """Test cases for extract_code_blocks function"""
 
@@ -52,7 +51,7 @@ class TestCreateAppFile:
         file_path = str(tmp_path)
         code = "import streamlit\nst.title('Test')"
         
-        success, message = create_app_file(file_path, code)
+        success, app_path, message = create_app_file(file_path, code)
         
         assert success is True
         assert "Successfully created" in message
@@ -70,7 +69,7 @@ class TestCreateAppFile:
         file_path = "/nonexistent/path/that/should/fail"
         code = "import streamlit"
         
-        success, message = create_app_file(file_path, code)
+        success, app_path, message = create_app_file(file_path, code)
         
         assert success is False
         assert "Error creating file" in message
@@ -81,7 +80,7 @@ class TestCreateAppFile:
         file_path = "/tmp/test"
         code = "import streamlit"
         
-        success, message = create_app_file(file_path, code)
+        success, app_path, message = create_app_file(file_path, code)
         
         assert success is False
         assert "Unexpected error" in message
@@ -91,7 +90,7 @@ class TestCreateAppFile:
         file_path = str(tmp_path)
         code = ""
         
-        success, message = create_app_file(file_path, code)
+        success, app_path, message = create_app_file(file_path, code)
         
         assert success is True
         assert "Successfully created" in message

mito_ai/tests/streamlit_preview/test_streamlit_preview_manager.py

@@ -0,0 +1,302 @@
+# Copyright (c) Saga Inc.
+# Distributed under the terms of the GNU Affero General Public License v3.0 License.
+
+import pytest
+import time
+import tempfile
+import os
+import shutil
+import subprocess
+import threading
+import requests
+from unittest.mock import Mock, patch, MagicMock
+from typing import Any
+
+from mito_ai.streamlit_preview.manager import (
+    StreamlitPreviewManager, 
+    PreviewProcess, 
+    get_preview_manager
+)
+
+
+class TestStreamlitPreviewManager:
+    """Test cases for StreamlitPreviewManager."""
+    
+    @pytest.fixture
+    def manager(self):
+        """Create a fresh manager instance for each test."""
+        return StreamlitPreviewManager()
+    
+    @pytest.fixture
+    def sample_app_code(self):
+        """Sample streamlit app code for testing."""
+        return """
+import streamlit as st
+
+st.title("Test App")
+st.write("Hello, World!")
+"""
+    
+    def test_init(self, manager):
+        """Test manager initialization."""
+        assert manager._previews == {}
+        assert isinstance(manager._lock, type(threading.Lock()))
+        assert manager.log is not None
+    
+    def test_get_free_port(self, manager):
+        """Test getting a free port."""
+        port = manager.get_free_port()
+        assert isinstance(port, int)
+        assert port > 0
+        assert port < 65536
+        
+        # Test that we get different ports
+        port2 = manager.get_free_port()
+        assert port != port2
+    
+    @pytest.mark.parametrize("app_code,preview_id,expected_success", [
+        ("import streamlit as st\nst.write('Hello')", "test_preview", True),
+        ("", "empty_preview", True),
+        ("import streamlit as st\n" * 1000 + "st.write('Large app')", "large_preview", True),
+    ])
+    def test_start_streamlit_preview_success_cases(self, manager, app_code, preview_id, expected_success):
+        """Test successful streamlit preview start with different app codes."""
+        with patch('subprocess.Popen') as mock_popen, \
+             patch('requests.get') as mock_requests_get, \
+             patch('tempfile.mkdtemp') as mock_mkdtemp:
+            
+            # Setup mocks
+            app_directory = "/tmp/test_dir"
+            mock_mkdtemp.return_value = app_directory
+            mock_proc = Mock()
+            mock_proc.terminate.return_value = None
+            mock_proc.wait.return_value = None
+            mock_popen.return_value = mock_proc
+            
+            mock_response = Mock()
+            mock_response.status_code = 200
+            mock_requests_get.return_value = mock_response
+            
+            # Test
+            success, message, 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)
+    
+    @pytest.mark.parametrize("exception_type,expected_message", [
+        (Exception("Temp dir creation failed"), "failed to start preview"),
+        (OSError("Permission denied"), "failed to start preview"),
+        (ValueError("Invalid argument"), "failed to start preview"),
+    ])
+    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):
+            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
+    
+    @pytest.mark.parametrize("preview_id,expected_result", [
+        ("existing_preview", True),
+        ("non_existent", False),
+    ])
+    def test_stop_preview_scenarios(self, manager, sample_app_code, preview_id, expected_result):
+        """Test stopping previews with different scenarios."""
+        if expected_result:
+            # Start a preview first
+            with patch('subprocess.Popen') as mock_popen, \
+                 patch('requests.get') as mock_requests_get, \
+                 patch('tempfile.mkdtemp') as mock_mkdtemp, \
+                 patch('builtins.open', create=True) as mock_open, \
+                 patch('os.path.exists') as mock_exists:
+                
+                app_directory = "/tmp/test_dir"
+                mock_mkdtemp.return_value = app_directory
+                mock_proc = Mock()
+                mock_proc.terminate.return_value = None
+                mock_proc.wait.return_value = None
+                mock_popen.return_value = mock_proc
+                
+                mock_response = Mock()
+                mock_response.status_code = 200
+                mock_requests_get.return_value = mock_response
+                
+                # Mock file operations
+                mock_file = Mock()
+                mock_open.return_value.__enter__.return_value = mock_file
+                mock_exists.return_value = True
+                
+                manager.start_streamlit_preview(app_directory, preview_id)
+    
+    @pytest.mark.parametrize("process_behavior,expected_kill_called", [
+        (subprocess.TimeoutExpired("cmd", 5), True),
+        (None, False),  # Normal termination
+    ])
+    def test_stop_preview_process_behaviors(self, manager, sample_app_code, process_behavior, expected_kill_called):
+        """Test stopping preview with different process behaviors."""
+        with patch('subprocess.Popen') as mock_popen, \
+             patch('requests.get') as mock_requests_get, \
+             patch('tempfile.mkdtemp') as mock_mkdtemp, \
+             patch('builtins.open', create=True) as mock_open, \
+             patch('os.path.exists') as mock_exists:
+            
+            # Setup mocks for start
+            app_directory = "/tmp/test_dir"
+            mock_mkdtemp.return_value = app_directory
+            
+            mock_proc = Mock()
+            mock_proc.terminate.return_value = None
+            mock_proc.wait.return_value = None
+            mock_popen.return_value = mock_proc
+            
+            mock_response = Mock()
+            mock_response.status_code = 200
+            mock_requests_get.return_value = mock_response
+            
+            # Mock file operations
+            mock_file = Mock()
+            mock_open.return_value.__enter__.return_value = mock_file
+            mock_exists.return_value = True
+            
+            # Start a preview
+            manager.start_streamlit_preview(app_directory, "test_preview")
+            
+            # Setup process behavior for stop
+            if process_behavior:
+                # Configure the mock to raise the exception when called with timeout
+                def wait_with_timeout(*args, **kwargs):
+                    if 'timeout' in kwargs:
+                        raise process_behavior
+                    return None
+                mock_proc.wait.side_effect = wait_with_timeout
+    
+    @pytest.mark.parametrize("preview_id,expected_found", [
+        ("existing_preview", True),
+        ("non_existent", False),
+    ])
+    def test_get_preview_scenarios(self, manager, sample_app_code, preview_id, expected_found):
+        """Test getting previews with different scenarios."""
+        if expected_found:
+            # Start a preview first
+            with patch('subprocess.Popen') as mock_popen, \
+                 patch('requests.get') as mock_requests_get, \
+                 patch('tempfile.mkdtemp') as mock_mkdtemp, \
+                 patch('builtins.open', create=True) as mock_open, \
+                 patch('os.path.exists') as mock_exists:
+                
+                mock_mkdtemp.return_value = "/tmp/test_dir"
+                mock_proc = Mock()
+                mock_proc.terminate.return_value = None
+                mock_proc.wait.return_value = None
+                mock_popen.return_value = mock_proc
+                
+                mock_response = Mock()
+                mock_response.status_code = 200
+                mock_requests_get.return_value = mock_response
+                
+                # Mock file operations
+                mock_file = Mock()
+                mock_open.return_value.__enter__.return_value = mock_file
+                mock_exists.return_value = True
+                
+                manager.start_streamlit_preview("/tmp/test_dir", preview_id)
+        
+        preview = manager.get_preview(preview_id)
+        
+        if expected_found:
+            assert preview is not None
+            assert isinstance(preview, PreviewProcess)
+            assert preview.port > 0
+            
+            # Cleanup
+            manager.stop_preview(preview_id)
+        else:
+            assert preview is None
+    
+    def test_preview_process_dataclass(self):
+        """Test PreviewProcess dataclass."""
+        proc = Mock()
+        port = 8080
+        
+        preview = PreviewProcess(
+            proc=proc,
+            port=port
+        )
+        
+        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."""
+        preview_ids = [f"preview_{i}" for i in range(num_previews)]
+        ports = []
+        
+        with patch('subprocess.Popen') as mock_popen, \
+             patch('requests.get') as mock_requests_get, \
+             patch('tempfile.mkdtemp') as mock_mkdtemp, \
+             patch('builtins.open', create=True) as mock_open, \
+             patch('os.path.exists') as mock_exists:
+            
+            # Setup mocks
+            mock_mkdtemp.return_value = "/tmp/test_dir"
+            mock_proc = Mock()
+            mock_proc.terminate.return_value = None
+            mock_proc.wait.return_value = None
+            mock_popen.return_value = mock_proc
+            
+            mock_response = Mock()
+            mock_response.status_code = 200
+            mock_requests_get.return_value = mock_response
+            
+            # Mock file operations
+            mock_file = Mock()
+            mock_open.return_value.__enter__.return_value = mock_file
+            mock_exists.return_value = True
+            
+            # Start multiple previews
+            for preview_id in preview_ids:
+                success, _, port = manager.start_streamlit_preview("/tmp/test_dir", preview_id)
+                assert success is True
+                ports.append(port)
+            
+            # Assertions
+            assert len(set(ports)) == num_previews  # All ports should be different
+            
+            # Check all previews exist
+            for preview_id in preview_ids:
+                assert manager.get_preview(preview_id) is not None
+            
+            # Stop all previews
+            for preview_id in preview_ids:
+                assert manager.stop_preview(preview_id)
+            
+            # Verify they're gone
+            for preview_id in preview_ids:
+                assert manager.get_preview(preview_id) is None
+    

mito_ai/utils/telemetry_utils.py

@@ -174,7 +174,6 @@ def log(
     If telemetry is not turned off and we are not running tests,
     we log the ai event
     """
-
     final_params: Dict[str, Any] = params or {}
     
     # Then, make sure to add the user email
@@ -355,4 +354,32 @@ def log_db_connection_error(connection_type: str, error_message: str) -> None:
             "error_message": error_message,
         }
     )
+    
+#################################
+# Streamlit Conversion
+#################################
 
+def log_streamlit_app_creation_success(key_type: Literal['mito_server_key', 'user_key'], message_type: MessageType) -> None:
+    log(
+        "mito_ai_streamlit_app_creation_success", 
+        key_type=key_type
+    )
+    
+def log_streamlit_app_creation_retry(key_type: Literal['mito_server_key', 'user_key'], message_type: MessageType, error: str) -> None:
+    log(
+        "mito_ai_streamlit_app_creation_retry", 
+        params={
+            "error_message": error,
+        },
+        key_type=key_type
+    )
+    
+def log_streamlit_app_creation_error(key_type: Literal['mito_server_key', 'user_key'], message_type: MessageType, error: str) -> None:
+    log(
+        "mito_ai_streamlit_app_creation_error", 
+        params={
+            "error_message": error,
+        },
+        key_type=key_type
+    )
+    

{mito_ai-0.1.37.data → mito_ai-0.1.38.data}/data/share/jupyter/labextensions/mito_ai/build_log.json

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

{mito_ai-0.1.37.data → mito_ai-0.1.38.data}/data/share/jupyter/labextensions/mito_ai/package.json

@@ -1,6 +1,6 @@
 {
   "name": "mito_ai",
-  "version": "0.1.37",
+  "version": "0.1.38",
   "description": "AI chat for JupyterLab",
   "keywords": [
     "jupyter",
@@ -138,7 +138,7 @@
     "outputDir": "mito_ai/labextension",
     "schemaDir": "schema",
     "_build": {
-      "load": "static/remoteEntry.93ecc9bc0edba61535cc.js",
+      "load": "static/remoteEntry.bcce4ea34631acf6dbbe.js",
       "extension": "./extension",
       "style": "./style"
     }

{mito_ai-0.1.37.data → mito_ai-0.1.38.data}/data/share/jupyter/labextensions/mito_ai/schemas/mito_ai/package.json.orig

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

{mito_ai-0.1.37.data → mito_ai-0.1.38.data}/data/share/jupyter/labextensions/mito_ai/schemas/mito_ai/toolbar-buttons.json

@@ -27,10 +27,15 @@
         }
       ],
       "Notebook": [
+        {
+          "name": "preview-app",
+          "command": "toolbar-button:preview-as-streamlit",
+          "rank": 1000
+        },
         {
           "name": "convert-to-streamlit",
           "command": "toolbar-button:convert-to-streamlit",
-          "rank": 1000
+          "rank": 1001
         }
       ]
     }