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

mito_ai/streamlit_preview/handlers.py

@@ -1,18 +1,17 @@
 # Copyright (c) Saga Inc.
 # Distributed under the terms of the GNU Affero General Public License v3.0 License.
 
-import os
-import tempfile
 import uuid
-from mito_ai.streamlit_conversion.streamlit_utils import get_app_path
-from mito_ai.streamlit_preview.utils import ensure_app_exists, validate_request_body
+from mito_ai.streamlit_preview.utils import validate_request_body
 import tornado
 from jupyter_server.base.handlers import APIHandler
+from mito_ai.streamlit_preview.manager import StreamlitPreviewManager
+from mito_ai.path_utils import get_absolute_notebook_dir_path, get_absolute_notebook_path, get_absolute_app_path, does_app_path_exist, get_app_file_name
+from mito_ai.utils.telemetry_utils import log_streamlit_app_conversion_error, log_streamlit_app_preview_failure, log_streamlit_app_preview_success
+from mito_ai.completions.models import MessageType
+from mito_ai.utils.error_classes import StreamlitConversionError, StreamlitPreviewError
 from mito_ai.streamlit_conversion.streamlit_agent_handler import streamlit_handler
-from mito_ai.streamlit_preview.manager import get_preview_manager
-from mito_ai.utils.create import initialize_user
-from typing import Tuple, Optional
-
+import traceback
 
 
 class StreamlitPreviewHandler(APIHandler):
@@ -20,61 +19,7 @@ class StreamlitPreviewHandler(APIHandler):
 
     def initialize(self) -> None:
         """Initialize the handler."""
-        self.preview_manager = get_preview_manager()
-
-    def _resolve_notebook_path(self, notebook_path: str) -> str:
-        """
-        Resolve the notebook path to an absolute path that can be found by the backend.
-
-        This method handles path resolution issues that can occur in different environments:
-
-        1. **Test Environment**: Playwright tests create temporary directories with complex
-           paths like 'mitoai_ui_tests-app_builde-ab3a5-n-Test-Preview-as-Streamlit-chromium/'
-           that the backend can't directly access.
-
-        2. **JupyterHub/Cloud Deployments**: In cloud environments, users may have notebooks
-           in subdirectories that aren't immediately accessible from the server root.
-
-        3. **Docker Containers**: When running in containers, the working directory and
-           file paths may not align with what the frontend reports.
-
-        4. **Multi-user Environments**: In enterprise deployments, users may have notebooks
-           in user-specific directories that require path resolution.
-
-        The method tries multiple strategies:
-        1. If the path is already absolute, return it as-is
-        2. Try to resolve relative to the Jupyter server's root directory
-        3. Search recursively through subdirectories for a file with the same name
-        4. Return the original path if not found (will cause a clear error message)
-
-        Args:
-            notebook_path (str): The notebook path from the frontend (may be relative or absolute)
-
-        Returns:
-            str: The resolved absolute path to the notebook file
-        """
-        # If the path is already absolute, return it
-        if os.path.isabs(notebook_path):
-            return notebook_path
-
-        # Get the Jupyter server's root directory
-        server_root = self.settings.get("server_root_dir", os.getcwd())
-
-        # Try to find the notebook file in the server root
-        resolved_path = os.path.join(server_root, notebook_path)
-        if os.path.exists(resolved_path):
-            return resolved_path
-
-        # If not found, try to find it in subdirectories
-        # This handles cases where the notebook is in a subdirectory that the frontend
-        # doesn't know about, or where the path structure differs between frontend and backend
-        for root, dirs, files in os.walk(server_root):
-            if os.path.basename(notebook_path) in files:
-                return os.path.join(root, os.path.basename(notebook_path))
-
-        # If still not found, return the original path (will cause a clear error)
-        # This ensures we get a meaningful error message rather than a generic "file not found"
-        return os.path.join(os.getcwd(), notebook_path)
+        self.preview_manager = StreamlitPreviewManager()
 
     @tornado.web.authenticated
     async def post(self) -> None:
@@ -82,42 +27,66 @@ class StreamlitPreviewHandler(APIHandler):
         try:
             # Parse and validate request
             body = self.get_json_body()
-            is_valid, error_msg, notebook_path, force_recreate, edit_prompt = validate_request_body(body)
-            if not is_valid or not notebook_path:
-                self.set_status(400)
-                self.finish({"error": error_msg})
-                return
+            notebook_path, notebook_id, force_recreate, edit_prompt = validate_request_body(body)
 
             # Ensure app exists
-            resolved_notebook_path = self._resolve_notebook_path(notebook_path)
+            absolute_notebook_path = get_absolute_notebook_path(notebook_path)
+            absolute_notebook_dir_path = get_absolute_notebook_dir_path(absolute_notebook_path)
+            app_file_name = get_app_file_name(notebook_id)
+            absolute_app_path = get_absolute_app_path(absolute_notebook_dir_path, app_file_name)
+            app_path_exists = does_app_path_exist(absolute_app_path)
 
-            success, error_msg = await ensure_app_exists(resolved_notebook_path, force_recreate, edit_prompt)
+            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")
 
-            if not success:
-                self.set_status(500)
-                self.finish({"error": error_msg})
-                return
+                await streamlit_handler(absolute_notebook_path, app_file_name, edit_prompt)
 
             # Start preview
             # TODO: There's a bug here where when the user rebuilds and already running app. Instead of 
             # creating a new process, we should update the existing process. The app displayed to the user 
             # does update, but that's just because of hot reloading when we overwrite the app.py file. 
             preview_id = str(uuid.uuid4())
-            resolved_app_directory = os.path.dirname(resolved_notebook_path)
-            success, message, port = self.preview_manager.start_streamlit_preview(resolved_app_directory, preview_id)
-
-            if not success:
-                self.set_status(500)
-                self.finish({"error": f"Failed to start preview: {message}"})
-                return
+            port = self.preview_manager.start_streamlit_preview(absolute_notebook_dir_path, app_file_name, preview_id)
 
             # Return success response
-            self.finish({"id": preview_id, "port": port, "url": f"http://localhost:{port}"})
-
+            self.finish({
+                "type": 'success',
+                "id": preview_id, 
+                "port": port, 
+                "url": f"http://localhost:{port}"
+            })
+            log_streamlit_app_preview_success('mito_server_key', MessageType.STREAMLIT_CONVERSION, edit_prompt)
+
+        except StreamlitConversionError as e:
+            print(e)
+            self.set_status(e.error_code)
+            error_message = str(e)
+            formatted_traceback = traceback.format_exc()
+            self.finish({"error": error_message})
+            log_streamlit_app_conversion_error(
+                'mito_server_key', 
+                MessageType.STREAMLIT_CONVERSION, 
+                error_message, 
+                formatted_traceback,
+                edit_prompt,
+            )
+        except StreamlitPreviewError as e:
+            print(e)
+            error_message = str(e)
+            formatted_traceback = traceback.format_exc()
+            self.set_status(e.error_code)
+            self.finish({"error": error_message})
+            log_streamlit_app_preview_failure('mito_server_key', MessageType.STREAMLIT_CONVERSION, error_message, formatted_traceback, edit_prompt)
         except Exception as e:
-            print(f"Error in streamlit preview handler: {e}")
+            print(f"Exception in streamlit preview handler: {e}")
             self.set_status(500)
-            self.finish({"error": str(e)})
+            error_message = str(e)
+            formatted_traceback = traceback.format_exc()
+            self.finish({"error": error_message})
+            log_streamlit_app_preview_failure('mito_server_key', MessageType.STREAMLIT_CONVERSION, error_message, formatted_traceback, "")
 
     @tornado.web.authenticated
     def delete(self, preview_id: str) -> None:

mito_ai/streamlit_preview/manager.py

@@ -1,16 +1,16 @@
 # 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.path_utils import AbsoluteNotebookDirPath, AppFileName
+from mito_ai.utils.error_classes import StreamlitPreviewError
 
 
 @dataclass
@@ -37,7 +37,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: AbsoluteNotebookDirPath, app_file_name: AppFileName, preview_id: str) -> int:
         """Start a streamlit preview process.
         
         Args:
@@ -47,6 +47,7 @@ class StreamlitPreviewManager:
         Returns:
             Tuple of (success, message, port)
         """
+        
         try:
             
             # Get free port
@@ -54,12 +55,12 @@ class StreamlitPreviewManager:
             
             # Start streamlit process
             cmd = [
-                "streamlit", "run", 'app.py', # Since we run this command from the app_directory, we always just run app.py 
+                "streamlit", "run", app_file_name,
                 "--server.port", str(port),
                 "--server.headless", "true",
                 "--server.address", "localhost",
                 "--server.enableXsrfProtection", "false",
-                "--server.runOnSave", "true",  # auto-reload when app.py is saved
+                "--server.runOnSave", "true",  # auto-reload when app is saved
                 "--logger.level", "error"
             ]
             
@@ -80,7 +81,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 +91,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 +107,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 +123,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 +150,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,28 @@
 # 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, 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)
+    
+    notebook_id = body.get("notebook_id")
+    if not notebook_id:
+        raise StreamlitPreviewError("Missing notebook_id 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, notebook_id, force_recreate, edit_prompt

mito_ai/tests/deploy_app/test_app_deploy_utils.py

@@ -4,6 +4,7 @@
 import zipfile
 import logging
 from mito_ai.app_deploy.app_deploy_utils import add_files_to_zip
+from mito_ai.path_utils import AbsoluteNotebookDirPath
 
 class TestAddFilesToZip:
     """Test cases for add_files_to_zip helper function"""
@@ -17,13 +18,30 @@ class TestAddFilesToZip:
         f2.write_text("file2 content")
 
         zip_path = tmp_path / "test.zip"
-        add_files_to_zip(str(zip_path), str(tmp_path), ["file1.txt", "file2.txt"])
+        add_files_to_zip(str(zip_path), AbsoluteNotebookDirPath(str(tmp_path)), ["file1.txt", "file2.txt"], 'test-app-file-name.py')
 
         with zipfile.ZipFile(zip_path, "r") as zf:
             names = zf.namelist()
             assert "file1.txt" in names
             assert "file2.txt" in names
             assert len(names) == 2
+            
+    def test_renames_app_file(self, tmp_path):
+        """Ensure individual files are added correctly to the zip"""
+        # Create files
+        f1 = tmp_path / "original-file-name.py"
+        f1.write_text("file1 content")
+        f2 = tmp_path / "file2.txt"
+        f2.write_text("file2 content")
+
+        zip_path = tmp_path / "test.zip"
+        add_files_to_zip(str(zip_path), AbsoluteNotebookDirPath(str(tmp_path)), ["original-file-name.py", "file2.txt"], 'original-file-name.py')
+
+        with zipfile.ZipFile(zip_path, "r") as zf:
+            names = zf.namelist()
+            assert "app.py" in names
+            assert "file2.txt" in names
+            assert len(names) == 2
 
     def test_directories_added_recursively(self, tmp_path):
         """Ensure directories are added recursively with correct relative paths"""
@@ -35,7 +53,7 @@ class TestAddFilesToZip:
         (subfolder / "nested2.txt").write_text("nested2 content")
 
         zip_path = tmp_path / "test.zip"
-        add_files_to_zip(str(zip_path), str(tmp_path), ["folder"])
+        add_files_to_zip(str(zip_path), AbsoluteNotebookDirPath(str(tmp_path)), ["folder"], 'test-app.py')
 
         with zipfile.ZipFile(zip_path, "r") as zf:
             names = zf.namelist()
@@ -46,7 +64,7 @@ class TestAddFilesToZip:
         """Ensure missing files do not break the function and warning is logged"""
         caplog.set_level(logging.WARNING)
         zip_path = tmp_path / "test.zip"
-        add_files_to_zip(str(zip_path), str(tmp_path), ["does_not_exist.txt"], logger=logging.getLogger())
+        add_files_to_zip(str(zip_path), AbsoluteNotebookDirPath(str(tmp_path)), ["does_not_exist.txt"], 'test-app.py', logger=logging.getLogger())
 
         # Zip should exist but be empty
         with zipfile.ZipFile(zip_path, "r") as zf:
@@ -63,7 +81,7 @@ class TestAddFilesToZip:
         (folder / "nested.txt").write_text("nested content")
 
         zip_path = tmp_path / "test.zip"
-        add_files_to_zip(str(zip_path), str(tmp_path), ["file.txt", "folder"])
+        add_files_to_zip(str(zip_path), AbsoluteNotebookDirPath(str(tmp_path)), ["file.txt", "folder"], 'test-app.py')
 
         with zipfile.ZipFile(zip_path, "r") as zf:
             names = zf.namelist()

mito_ai/tests/message_history/test_message_history_utils.py

@@ -106,12 +106,15 @@ PROMPT_BUILDER_TEST_CASES = [
             AgentExecutionMetadata(
                 variables=TEST_VARIABLES,
                 files=TEST_FILES,
+                notebookPath='/test-notebook-path.ipynb',
+                notebookID='test-notebook-id',
                 aiOptimizedCells=[
                     AIOptimizedCell(cell_type="code", id="cell1", code=TEST_CODE)
                 ],
                 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\")")])
+
+
+