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

mito_ai/streamlit_preview/utils.py

@@ -2,44 +2,41 @@
 # 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.path_utils import AbsoluteNotebookPath, does_app_path_exists, get_absolute_app_path, get_absolute_notebook_dir_path
+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, ""
+        raise StreamlitPreviewError("edit_prompt must be a string", 400)
 
-    return True, "", notebook_path, force_recreate, edit_prompt
+    return 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]:
+async def ensure_app_exists(absolute_notebook_path: AbsoluteNotebookPath, force_recreate: bool = False, edit_prompt: str = "") -> None:
     """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))
+
+    absolute_notebook_dir_path = get_absolute_notebook_dir_path(absolute_notebook_path)
+    absolute_app_path = get_absolute_app_path(absolute_notebook_dir_path)
+    app_path_exists = does_app_path_exists(absolute_app_path)
     
-    if app_path is None or force_recreate:
-        if app_path is None:
+    if not app_path_exists or force_recreate:
+        if not app_path_exists:
             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}"
-
-    return True, ""
+        await streamlit_handler(absolute_notebook_path, edit_prompt)

mito_ai/tests/streamlit_conversion/test_apply_search_replace.py

@@ -0,0 +1,226 @@
+# 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")"""
+    )
+])
+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\")")])
+
+
+def test_apply_search_replace_multiple_matches():
+    """Test that ValueError is raised when search text is found multiple times."""
+    with pytest.raises(StreamlitConversionError, match="Search text found 2 times"):
+        apply_search_replace("st.write(\"Hello\")\nst.write(\"Hello\")", [("st.write(\"Hello\")", "st.write(\"Hi\")")])

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,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_success(self, mock_clean_directory, mock_create_file, mock_validator, mock_generate_code, mock_parse):
+    async def test_streamlit_handler_success(self, 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"]}]}]
@@ -153,26 +150,19 @@ class TestStreamlitHandler:
         # 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
-        
         # 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')
@@ -192,11 +182,9 @@ class TestStreamlitHandler:
         # Mock validation (always errors) - Return list of errors as expected by validate_app
         mock_validator.return_value = (True, ["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)
         assert mock_correct_error.call_count == 5
@@ -206,8 +194,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": []}]
@@ -219,35 +206,27 @@ class TestStreamlitHandler:
         # Mock validation (no errors)
         mock_validator.return_value = (False, "")
         
-        # Mock file creation failure
-        mock_create_file.return_value = (False, None, "Permission denied")
-        
-        # 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, match="Permission denied"):
+            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 +235,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"))
 
 
 

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

@@ -11,6 +11,7 @@ from mito_ai.streamlit_conversion.validate_streamlit_app import (
     validate_app
 )
 import pytest
+from mito_ai.path_utils import AbsoluteNotebookPath
 
 
 class TestGetSyntaxError:
@@ -56,7 +57,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,7 +88,7 @@ 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:
@@ -99,7 +102,7 @@ class TestValidateApp:
     ])
     def test_validate_app(self, app_code, expected_has_validation_error, expected_error_message):
         """Test the validate_app function"""
-        has_validation_error, errors = validate_app(app_code, '/app.py')
+        has_validation_error, errors = validate_app(app_code, AbsoluteNotebookPath('/notebook.ipynb'))
         
         assert has_validation_error == expected_has_validation_error
         assert expected_error_message in str(errors)