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

mito_ai/streamlit_preview/manager.py

@@ -1,16 +1,15 @@
 # Copyright (c) Saga Inc.
 # Distributed under the terms of the GNU Affero General Public License v3.0 License.
 
-import os
 import socket
 import subprocess
-import tempfile
 import time
 import threading
 import requests
 from typing import Dict, Optional, Tuple
 from dataclasses import dataclass
 from mito_ai.logger import get_logger
+from mito_ai.utils.error_classes import StreamlitPreviewError
 
 
 @dataclass
@@ -37,7 +36,7 @@ class StreamlitPreviewManager:
         
         return port
     
-    def start_streamlit_preview(self, app_directory: str, preview_id: str) -> Tuple[bool, str, Optional[int]]:
+    def start_streamlit_preview(self, app_directory: str, preview_id: str) -> int:
         """Start a streamlit preview process.
         
         Args:
@@ -80,7 +79,7 @@ class StreamlitPreviewManager:
             if not ready:
                 proc.terminate()
                 proc.wait()
-                return False, "Streamlit app failed to start", None
+                raise StreamlitPreviewError("Streamlit app failed to start as app is not ready", 500)
             
             # Register the process
             with self._lock:
@@ -90,11 +89,11 @@ class StreamlitPreviewManager:
                 )
             
             self.log.info(f"Started streamlit preview {preview_id} on port {port}")
-            return True, "Preview started successfully", port
+            return port
             
         except Exception as e:
             self.log.error(f"Error starting streamlit preview: {e}")
-            return False, f"Failed to start preview: {str(e)}", None
+            raise StreamlitPreviewError(f"Failed to start preview: {str(e)}", 500)
     
     def _wait_for_app_ready(self, port: int, timeout: int = 30) -> bool:
         """Wait for streamlit app to be ready on the given port."""
@@ -106,7 +105,7 @@ class StreamlitPreviewManager:
                 if response.status_code == 200:
                     return True
             except requests.exceptions.RequestException as e:
-                print(f"Error waiting for app to be ready: {e}")
+                self.log.info(f"Waiting for app to be ready...")
                 pass
             
             time.sleep(1)
@@ -122,7 +121,7 @@ class StreamlitPreviewManager:
         Returns:
             True if stopped successfully, False if not found
         """
-        print(f"Stopping preview {preview_id}")
+        self.log.info(f"Stopping preview {preview_id}")
         with self._lock:
             if preview_id not in self._previews:
                 return False
@@ -149,11 +148,3 @@ class StreamlitPreviewManager:
         """Get a preview process by ID."""
         with self._lock:
             return self._previews.get(preview_id)
-
-# Global instance
-_preview_manager = StreamlitPreviewManager()
-
-
-def get_preview_manager() -> StreamlitPreviewManager:
-    """Get the global preview manager instance."""
-    return _preview_manager 

mito_ai/streamlit_preview/utils.py

@@ -2,44 +2,24 @@
 # Distributed under the terms of the GNU Affero General Public License v3.0 License.
 
 from typing import Tuple, Optional
-import os
-from mito_ai.streamlit_conversion.streamlit_utils import get_app_path
-from mito_ai.streamlit_conversion.streamlit_agent_handler import streamlit_handler
+from mito_ai.utils.error_classes import StreamlitPreviewError
 
 
-def validate_request_body(body: Optional[dict]) -> Tuple[bool, str, Optional[str], bool, str]:
+def validate_request_body(body: Optional[dict]) -> Tuple[str, bool, str]:
     """Validate the request body and extract notebook_path and force_recreate."""
     if body is None:
-        return False, "Invalid or missing JSON body", None, False, ""
+        raise StreamlitPreviewError("Invalid or missing JSON body", 400)
 
     notebook_path = body.get("notebook_path")
     if not notebook_path:
-        return False, "Missing notebook_path parameter", None, False, ""
+        raise StreamlitPreviewError("Missing notebook_path parameter", 400)
 
     force_recreate = body.get("force_recreate", False)
     if not isinstance(force_recreate, bool):
-        return False, "force_recreate must be a boolean", None, False, ""
-    
+        raise StreamlitPreviewError("force_recreate must be a boolean", 400)
+
     edit_prompt = body.get("edit_prompt", "")
     if not isinstance(edit_prompt, str):
-        return False, "edit_prompt must be a string", None, False, ""
-
-    return True, "", notebook_path, force_recreate, edit_prompt
-
-async def ensure_app_exists(resolved_notebook_path: str, force_recreate: bool = False, edit_prompt: str = "") -> Tuple[bool, str]:
-    """Ensure app.py exists, generating it if necessary or if force_recreate is True."""
-    # Check if the app already exists
-    app_path = get_app_path(os.path.dirname(resolved_notebook_path))
-    
-    if app_path is None or force_recreate:
-        if app_path is None:
-            print("[Mito AI] App path not found, generating streamlit code")
-        else:
-            print("[Mito AI] Force recreating streamlit app")
-        
-        success, app_path, message = await streamlit_handler(resolved_notebook_path, edit_prompt)
-
-        if not success or app_path is None:
-            return False, f"Failed to generate streamlit code: {message}"
+        raise StreamlitPreviewError("edit_prompt must be a string", 400)
 
-    return True, ""
+    return notebook_path, force_recreate, edit_prompt

mito_ai/tests/message_history/test_message_history_utils.py

@@ -112,6 +112,7 @@ PROMPT_BUILDER_TEST_CASES = [
                 input=TEST_INPUT,
                 promptType="agent:execution",
                 threadId=ThreadID("test-thread-id"),
+                activeCellId="cell1",
                 isChromeBrowser=True
             )
         ),

mito_ai/tests/streamlit_conversion/test_apply_search_replace.py

@@ -0,0 +1,240 @@
+# Copyright (c) Saga Inc.
+# Distributed under the terms of the GNU Affero General Public License v3.0 License.
+
+from mito_ai.utils.error_classes import StreamlitConversionError
+import pytest
+from mito_ai.streamlit_conversion.search_replace_utils import apply_search_replace
+
+
+@pytest.mark.parametrize("original_text,search_replace_pairs,expected_result", [
+    # Test case 1: Simple title change
+    (
+        """import streamlit as st
+
+st.markdown(\"\"\"
+    <style>
+        #MainMenu {visibility: hidden;}
+        .stAppDeployButton {display:none;}
+        footer {visibility: hidden;}
+        .stMainBlockContainer {padding: 2rem 1rem 2rem 1rem;}
+    </style>
+\"\"\", unsafe_allow_html=True)
+
+st.title("Simple Calculation")
+
+x = 5
+y = 10
+result = x + y
+
+st.write(f"x = {x}")
+st.write(f"y = {y}")
+st.write(f"x + y = {result}")""",
+        [("st.title(\"Simple Calculation\")", "st.title(\"Math Examples\")")],
+        """import streamlit as st
+
+st.markdown(\"\"\"
+    <style>
+        #MainMenu {visibility: hidden;}
+        .stAppDeployButton {display:none;}
+        footer {visibility: hidden;}
+        .stMainBlockContainer {padding: 2rem 1rem 2rem 1rem;}
+    </style>
+\"\"\", unsafe_allow_html=True)
+
+st.title("Math Examples")
+
+x = 5
+y = 10
+result = x + y
+
+st.write(f"x = {x}")
+st.write(f"y = {y}")
+st.write(f"x + y = {result}")"""
+    ),
+    
+    # Test case 2: Add new content
+    (
+        """import streamlit as st
+
+st.title("My App")""",
+        [("st.title(\"My App\")", """st.title("My App")
+st.header("Welcome")
+st.write("This is a test app")""")],
+        """import streamlit as st
+
+st.title("My App")
+st.header("Welcome")
+st.write("This is a test app")"""
+    ),
+    
+    # Test case 3: Remove lines
+    (
+        """import streamlit as st
+
+st.header("Welcome")
+st.title("My App")
+st.write("This is a test app")""",
+        [("""st.header("Welcome")
+st.title("My App")
+st.write("This is a test app")""", "st.title(\"My App\")")],
+        """import streamlit as st
+
+st.title("My App")"""
+    ),
+    
+    # Test case 4: Multiple replacements
+    (
+        """import streamlit as st
+
+st.title("Old Title")
+x = 5
+y = 10
+st.write("Old message")""",
+        [
+            ("st.title(\"Old Title\")", "st.title(\"New Title\")"),
+            ("st.write(\"Old message\")", "st.write(\"New message\")")
+        ],
+        """import streamlit as st
+
+st.title("New Title")
+x = 5
+y = 10
+st.write("New message")"""
+    ),
+    
+    # Test case 5: Empty search/replace pairs
+    (
+        """import streamlit as st
+
+st.title("My App")""",
+        [],
+        """import streamlit as st
+
+st.title("My App")"""
+    ),
+    
+    # Test case 6: Complex replacement with context
+    (
+        """import streamlit as st
+
+# This is a comment
+st.title("Old Title")
+# Another comment
+x = 5
+y = 10
+# Final comment""",
+        [("""# This is a comment
+st.title("Old Title")
+# Another comment""", """# This is a comment
+st.title("New Title")
+# Another comment""")],
+        """import streamlit as st
+
+# This is a comment
+st.title("New Title")
+# Another comment
+x = 5
+y = 10
+# Final comment"""
+    ),
+    
+    # Test case 7: Replace multiple consecutive lines
+    (
+        """import streamlit as st
+
+st.title("My App")
+st.write("Line 1")
+st.write("Line 2")
+st.write("Line 3")
+
+x = 5""",
+        [("""st.write("Line 1")
+st.write("Line 2")
+st.write("Line 3")""", "st.write(\"New content\")")],
+        """import streamlit as st
+
+st.title("My App")
+st.write("New content")
+
+x = 5"""
+    ),
+    
+    # Test case 8: Add lines at the beginning
+    (
+        """import streamlit as st
+
+st.title("My App")""",
+        [("import streamlit as st", """import pandas as pd
+import streamlit as st""")],
+        """import pandas as pd
+import streamlit as st
+
+st.title("My App")"""
+    ),
+    
+    # Test case 9: Add lines at the end
+    (
+        """import streamlit as st
+
+st.title("My App")""",
+        [("st.title(\"My App\")", """st.title("My App")
+
+st.write("Footer content")
+st.write("More footer")""")],
+        """import streamlit as st
+
+st.title("My App")
+
+st.write("Footer content")
+st.write("More footer")"""
+    ),
+    
+    # Test case 10: Add emoji to streamlit app title
+    (
+        """import streamlit as st
+
+st.title("My App")
+st.write("Welcome to my application")""",
+        [("st.title(\"My App\")", "st.title(\"🚀 My App\")")],
+        """import streamlit as st
+
+st.title("🚀 My App")
+st.write("Welcome to my application")"""
+    ),
+    
+    # Test case 11: Only replace first occurrence when search text exists multiple times
+    (
+        """import streamlit as st
+
+st.write("Hello World")
+st.title("My App")
+st.write("Hello World")
+st.write("Another message")""",
+        [("st.write(\"Hello World\")", "st.write(\"Hi There\")")],
+        """import streamlit as st
+
+st.write("Hi There")
+st.title("My App")
+st.write("Hello World")
+st.write("Another message")"""
+    )
+])
+def test_apply_search_replace(original_text, search_replace_pairs, expected_result):
+    """Test the apply_search_replace function with various search/replace scenarios."""
+    result = apply_search_replace(original_text, search_replace_pairs)
+    
+    print(f"Original text: {repr(original_text)}")
+    print(f"Search/replace pairs: {search_replace_pairs}")
+    print(f"Expected result: {repr(expected_result)}")
+    print(f"Result: {repr(result)}")
+    
+    assert result == expected_result
+
+
+def test_apply_search_replace_search_not_found():
+    """Test that ValueError is raised when search text is not found."""
+    with pytest.raises(StreamlitConversionError, match="Search text not found"):
+        apply_search_replace("st.title(\"My App\")", [("st.title(\"Not Found\")", "st.title(\"New Title\")")])
+
+
+

mito_ai/tests/streamlit_conversion/test_streamlit_agent_handler.py

@@ -12,7 +12,7 @@ from mito_ai.streamlit_conversion.streamlit_agent_handler import (
     correct_error_in_generation,
     streamlit_handler
 )
-from mito_ai.streamlit_conversion.streamlit_utils import clean_directory_check
+from mito_ai.path_utils import AbsoluteAppPath, AbsoluteNotebookPath, get_absolute_app_path, get_absolute_notebook_dir_path, get_absolute_notebook_path
 
 # Add this line to enable async support
 pytest_plugins = ('pytest_asyncio',)
@@ -102,14 +102,12 @@ class TestCorrectErrorInGeneration:
     @patch('mito_ai.streamlit_conversion.agent_utils.stream_anthropic_completion_from_mito_server')
     async def test_correct_error_in_generation_success(self, mock_stream):
         """Test successful error correction"""
-        mock_response = """```unified_diff
---- a/app.py
-+++ b/app.py
-@@ -1,1 +1,1 @@
--import streamlit
--st.title('Test')
-+import streamlit
-+st.title('Fixed')
+        mock_response = """```search_replace
+>>>>>>> SEARCH
+st.title('Test')
+=======
+st.title('Fixed')
+<<<<<<< REPLACE
 ```"""
         async def mock_async_gen():
             for item in [mock_response]:
@@ -117,8 +115,8 @@ class TestCorrectErrorInGeneration:
 
         mock_stream.return_value = mock_async_gen()
 
-        result = await correct_error_in_generation("ImportError: No module named 'pandas'", "import streamlit\nst.title('Test')")
-
+        result = await correct_error_in_generation("ImportError: No module named 'pandas'", "import streamlit\nst.title('Test')\n")
+        
         expected_code = "import streamlit\nst.title('Fixed')\n"
         assert result == expected_code
 
@@ -140,8 +138,8 @@ class TestStreamlitHandler:
     @patch('mito_ai.streamlit_conversion.streamlit_agent_handler.generate_new_streamlit_code')
     @patch('mito_ai.streamlit_conversion.streamlit_agent_handler.validate_app')
     @patch('mito_ai.streamlit_conversion.streamlit_agent_handler.create_app_file')
-    @patch('mito_ai.streamlit_conversion.streamlit_agent_handler.clean_directory_check')
-    async def test_streamlit_handler_success(self, mock_clean_directory, mock_create_file, mock_validator, mock_generate_code, mock_parse):
+    @patch('mito_ai.streamlit_conversion.streamlit_agent_handler.log_streamlit_app_conversion_success')
+    async def test_streamlit_handler_success(self, mock_log_success, mock_create_file, mock_validator, mock_generate_code, mock_parse):
         """Test successful streamlit handler execution"""
         # Mock notebook parsing
         mock_notebook_data: List[dict] = [{"cells": [{"cell_type": "code", "source": ["import pandas"]}]}]
@@ -151,35 +149,29 @@ class TestStreamlitHandler:
         mock_generate_code.return_value = "import streamlit\nst.title('Test')"
         
         # Mock validation (no errors)
-        mock_validator.return_value = (False, "")
-        
-        # Mock file creation
-        mock_create_file.return_value = (True, "/path/to/app.py", "File created successfully")
-        
-        # Mock clean directory check (no-op)
-        mock_clean_directory.return_value = None
+        mock_validator.return_value = []
         
         # Use a relative path that will work cross-platform
-        notebook_path = "notebook.ipynb"
-        result = await streamlit_handler(notebook_path)
+        notebook_path = AbsoluteNotebookPath("absolute/path/to/notebook.ipynb")
         
-        assert result[0] is True
-        assert "File created successfully" in result[2]
+        # Construct the expected app path using the same method as the production code
+        app_directory = get_absolute_notebook_dir_path(notebook_path)
+        expected_app_path = get_absolute_app_path(app_directory)
+        await streamlit_handler(notebook_path)
         
         # Verify calls
         mock_parse.assert_called_once_with(notebook_path)
         mock_generate_code.assert_called_once_with(mock_notebook_data)
         mock_validator.assert_called_once_with("import streamlit\nst.title('Test')", notebook_path)
-        # get_app_directory converts relative paths to absolute, so expect the absolute path directory
-        expected_app_dir = os.path.dirname(os.path.abspath(notebook_path))
-        mock_create_file.assert_called_once_with(expected_app_dir, "import streamlit\nst.title('Test')")
+        mock_create_file.assert_called_once_with(expected_app_path, "import streamlit\nst.title('Test')")
 
     @pytest.mark.asyncio
     @patch('mito_ai.streamlit_conversion.streamlit_agent_handler.parse_jupyter_notebook_to_extract_required_content')
     @patch('mito_ai.streamlit_conversion.streamlit_agent_handler.generate_new_streamlit_code')
     @patch('mito_ai.streamlit_conversion.streamlit_agent_handler.correct_error_in_generation')
     @patch('mito_ai.streamlit_conversion.streamlit_agent_handler.validate_app')
-    async def test_streamlit_handler_max_retries_exceeded(self, mock_validator, mock_correct_error, mock_generate_code, mock_parse):
+    @patch('mito_ai.streamlit_conversion.streamlit_agent_handler.log_streamlit_app_validation_retry')
+    async def test_streamlit_handler_max_retries_exceeded(self, mock_log_retry, mock_validator, mock_correct_error, mock_generate_code, mock_parse):
         """Test streamlit handler when max retries are exceeded"""
         # Mock notebook parsing
         mock_notebook_data: List[dict] = [{"cells": []}]
@@ -189,16 +181,15 @@ class TestStreamlitHandler:
         mock_generate_code.return_value = "import streamlit\nst.title('Test')"
         mock_correct_error.return_value = "import streamlit\nst.title('Fixed')"
     
-        # Mock validation (always errors) - Return list of errors as expected by validate_app
-        mock_validator.return_value = (True, ["Persistent error"])
+        # Mock validation (always errors) - validate_app returns List[str]
+        mock_validator.return_value = ["Persistent error"]
     
-        result = await streamlit_handler("notebook.ipynb")
-        
-        # Verify the result indicates failure
-        assert result[0] is False
-        assert "Error generating streamlit code by agent" in result[2]
+        # Now it should raise an exception instead of returning a tuple
+        with pytest.raises(Exception):
+            await streamlit_handler(AbsoluteNotebookPath("notebook.ipynb"))
         
-        # Verify that error correction was called 5 times (max retries)
+        # Verify that error correction was called 5 times (once per error, 5 retries)
+        # Each retry processes 1 error, so 5 retries = 5 calls
         assert mock_correct_error.call_count == 5
 
     @pytest.mark.asyncio
@@ -206,8 +197,7 @@ class TestStreamlitHandler:
     @patch('mito_ai.streamlit_conversion.streamlit_agent_handler.generate_new_streamlit_code')
     @patch('mito_ai.streamlit_conversion.streamlit_agent_handler.validate_app')
     @patch('mito_ai.streamlit_conversion.streamlit_agent_handler.create_app_file')
-    @patch('mito_ai.streamlit_conversion.streamlit_agent_handler.clean_directory_check')
-    async def test_streamlit_handler_file_creation_failure(self, mock_clean_directory, mock_create_file, mock_validator, mock_generate_code, mock_parse):
+    async def test_streamlit_handler_file_creation_failure(self, mock_create_file, mock_validator, mock_generate_code, mock_parse):
         """Test streamlit handler when file creation fails"""
         # Mock notebook parsing
         mock_notebook_data: List[dict] = [{"cells": []}]
@@ -216,38 +206,30 @@ class TestStreamlitHandler:
         # Mock code generation
         mock_generate_code.return_value = "import streamlit\nst.title('Test')"
         
-        # Mock validation (no errors)
-        mock_validator.return_value = (False, "")
-        
-        # Mock file creation failure
-        mock_create_file.return_value = (False, None, "Permission denied")
+        # Mock validation (no errors) - validate_app returns List[str]
+        mock_validator.return_value = []
         
-        # Mock clean directory check (no-op)
-        mock_clean_directory.return_value = None
+        # Mock file creation failure - now it should raise an exception
+        mock_create_file.side_effect = Exception("Permission denied")
         
-        result = await streamlit_handler("notebook.ipynb")
-        
-        assert result[0] is False
-        assert "Permission denied" in result[2]
+        # Now it should raise an exception instead of returning a tuple
+        with pytest.raises(Exception):
+            await streamlit_handler(AbsoluteNotebookPath("notebook.ipynb"))
 
     @pytest.mark.asyncio
     @patch('mito_ai.streamlit_conversion.streamlit_agent_handler.parse_jupyter_notebook_to_extract_required_content')
-    @patch('mito_ai.streamlit_conversion.streamlit_agent_handler.clean_directory_check')
-    async def test_streamlit_handler_parse_notebook_exception(self, mock_clean_directory, mock_parse):
+    async def test_streamlit_handler_parse_notebook_exception(self, mock_parse):
         """Test streamlit handler when notebook parsing fails"""
-        # Mock clean directory check (no-op)
-        mock_clean_directory.return_value = None
         
         mock_parse.side_effect = FileNotFoundError("Notebook not found")
         
         with pytest.raises(FileNotFoundError, match="Notebook not found"):
-            await streamlit_handler("notebook.ipynb")
+            await streamlit_handler(AbsoluteNotebookPath("notebook.ipynb"))
 
     @pytest.mark.asyncio
     @patch('mito_ai.streamlit_conversion.streamlit_agent_handler.parse_jupyter_notebook_to_extract_required_content')
     @patch('mito_ai.streamlit_conversion.streamlit_agent_handler.generate_new_streamlit_code')
-    @patch('mito_ai.streamlit_conversion.streamlit_agent_handler.clean_directory_check')
-    async def test_streamlit_handler_generation_exception(self, mock_clean_directory, mock_generate_code, mock_parse):
+    async def test_streamlit_handler_generation_exception(self, mock_generate_code, mock_parse):
         """Test streamlit handler when code generation fails"""
         # Mock notebook parsing
         mock_notebook_data: List[dict] = [{"cells": []}]
@@ -256,11 +238,8 @@ class TestStreamlitHandler:
         # Mock code generation failure
         mock_generate_code.side_effect = Exception("Generation failed")
         
-        # Mock clean directory check (no-op)
-        mock_clean_directory.return_value = None
-        
         with pytest.raises(Exception, match="Generation failed"):
-            await streamlit_handler("notebook.ipynb")
+            await streamlit_handler(AbsoluteNotebookPath("notebook.ipynb"))