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

mito_ai/streamlit_conversion/search_replace_utils.py

@@ -0,0 +1,93 @@
+# Copyright (c) Saga Inc.
+# Distributed under the terms of the GNU Affero General Public License v3.0 License.
+
+import re
+from typing import List, Tuple
+
+from mito_ai.utils.error_classes import StreamlitConversionError
+
+
+def extract_search_replace_blocks(message_content: str) -> List[Tuple[str, str]]:
+    """
+    Extract all search_replace blocks from Claude's response.
+    
+    Returns:
+        List of tuples (search_text, replace_text) for each search/replace block
+    """
+    if "```search_replace" not in message_content:
+        return []
+    
+    pattern = r'```search_replace\n(.*?)```'
+    matches = re.findall(pattern, message_content, re.DOTALL)
+    
+    search_replace_pairs = []
+    for match in matches:
+        # Split by the separator
+        if "=======" not in match:
+            continue
+            
+        parts = match.split("=======", 1)
+        if len(parts) != 2:
+            continue
+            
+        search_part = parts[0]
+        replace_part = parts[1]
+        
+        # Extract search text (after SEARCH marker)
+        if ">>>>>>> SEARCH" in search_part:
+            search_text = search_part.split(">>>>>>> SEARCH", 1)[1].strip()
+        else:
+            continue
+            
+        # Extract replace text (before REPLACE marker)
+        if "<<<<<<< REPLACE" in replace_part:
+            replace_text = replace_part.split("<<<<<<< REPLACE", 1)[0].strip()
+        else:
+            continue
+            
+        search_replace_pairs.append((search_text, replace_text))
+    
+    return search_replace_pairs
+
+
+def apply_search_replace(text: str, search_replace_pairs: List[Tuple[str, str]]) -> str:
+    """
+    Apply search/replace operations to the given text.
+    
+    Parameters
+    ----------
+    text : str
+        The original file contents.
+    search_replace_pairs : List[Tuple[str, str]]
+        List of (search_text, replace_text) tuples to apply.
+        
+    Returns
+    -------
+    str
+        The updated contents after applying all search/replace operations.
+        
+    Raises
+    ------
+    ValueError
+        If a search text is not found or found multiple times.
+    """
+    if not search_replace_pairs:
+        return text
+    
+    result = text
+    
+    for search_text, replace_text in search_replace_pairs:
+        # Count occurrences of search text
+        count = result.count(search_text)
+        
+        if count == 0:
+            print("Search Text Not Found: ", repr(search_text))
+            raise StreamlitConversionError(f"Search text not found: {repr(search_text)}", error_code=500)
+        elif count > 1:
+            print("Search Text Found Multiple Times: ", repr(search_text))
+            raise StreamlitConversionError(f"Search text found {count} times (must be unique): {repr(search_text)}", error_code=500)
+        
+        # Perform the replacement
+        result = result.replace(search_text, replace_text)
+    
+    return result

mito_ai/streamlit_conversion/streamlit_agent_handler.py

@@ -4,24 +4,18 @@
 import os
 from anthropic.types import MessageParam
 from typing import List, Optional, Tuple, cast
-from mito_ai.streamlit_conversion.agent_utils import apply_patch_to_text, extract_todo_placeholders, fix_diff_headers, get_response_from_agent
+from mito_ai.streamlit_conversion.agent_utils import extract_todo_placeholders, get_response_from_agent
 from mito_ai.streamlit_conversion.prompts.streamlit_app_creation_prompt import get_streamlit_app_creation_prompt
 from mito_ai.streamlit_conversion.prompts.streamlit_error_correction_prompt import get_streamlit_error_correction_prompt
 from mito_ai.streamlit_conversion.prompts.streamlit_finish_todo_prompt import get_finish_todo_prompt
 from mito_ai.streamlit_conversion.prompts.update_existing_app_prompt import get_update_existing_app_prompt
 from mito_ai.streamlit_conversion.validate_streamlit_app import validate_app
-from mito_ai.streamlit_conversion.streamlit_utils import extract_code_blocks, create_app_file, extract_unified_diff_blocks, get_app_code_from_file, parse_jupyter_notebook_to_extract_required_content
+from mito_ai.streamlit_conversion.streamlit_utils import extract_code_blocks, create_app_file, get_app_code_from_file, parse_jupyter_notebook_to_extract_required_content
+from mito_ai.streamlit_conversion.search_replace_utils import extract_search_replace_blocks, apply_search_replace
 from mito_ai.completions.models import MessageType
-from mito_ai.utils.telemetry_utils import log_streamlit_app_creation_error, log_streamlit_app_creation_retry, log_streamlit_app_creation_success
-from mito_ai.streamlit_conversion.streamlit_utils import clean_directory_check
-
-def get_app_directory(notebook_path: str) -> str:
-    # Make sure the path is absolute if it is not already
-    absolute_notebook_path = os.path.abspath(notebook_path)
-    
-    # Get the directory of the notebook
-    app_directory = os.path.dirname(absolute_notebook_path)
-    return app_directory
+from mito_ai.utils.error_classes import StreamlitConversionError
+from mito_ai.utils.telemetry_utils import log_streamlit_app_validation_retry, log_streamlit_app_conversion_success
+from mito_ai.path_utils import AbsoluteNotebookPath, get_absolute_notebook_dir_path, get_absolute_app_path
 
 async def generate_new_streamlit_code(notebook: List[dict]) -> str:
     """Send a query to the agent, get its response and parse the code"""
@@ -57,10 +51,9 @@ async def generate_new_streamlit_code(notebook: List[dict]) -> str:
         ]
         todo_response = await get_response_from_agent(todo_messages)
         
-        # Apply the diff to the streamlit app
-        exctracted_diff = extract_unified_diff_blocks(todo_response)
-        fixed_diff = fix_diff_headers(exctracted_diff)
-        converted_code = apply_patch_to_text(converted_code, fixed_diff)
+        # Apply the search/replace to the streamlit app
+        search_replace_pairs = extract_search_replace_blocks(todo_response)
+        converted_code = apply_search_replace(converted_code, search_replace_pairs)
                 
     return converted_code
 
@@ -80,10 +73,12 @@ async def update_existing_streamlit_code(notebook: List[dict], streamlit_app_cod
     ]
     
     agent_response = await get_response_from_agent(messages)
-    exctracted_diff = extract_unified_diff_blocks(agent_response)
-    fixed_diff = fix_diff_headers(exctracted_diff)
-    print(fixed_diff)
-    converted_code = apply_patch_to_text(streamlit_app_code, fixed_diff)
+    print(f"[Mito AI Search/Replace Tool]:\n {agent_response}")
+
+    # Apply the search/replace to the streamlit app
+    search_replace_pairs = extract_search_replace_blocks(agent_response)
+    converted_code = apply_search_replace(streamlit_app_code, search_replace_pairs)
+    print(f"[Mito AI Search/Replace Tool]\nConverted code\n: {converted_code}")
     return converted_code
 
 
@@ -100,27 +95,26 @@ async def correct_error_in_generation(error: str, streamlit_app_code: str) -> st
     ]
     agent_response = await get_response_from_agent(messages)
     
-    # Apply the diff to the streamlit app
-    exctracted_diff = extract_unified_diff_blocks(agent_response)
-    fixed_diff = fix_diff_headers(exctracted_diff)
-    streamlit_app_code = apply_patch_to_text(streamlit_app_code, fixed_diff)
+    # Apply the search/replace to the streamlit app
+    search_replace_pairs = extract_search_replace_blocks(agent_response)
+    streamlit_app_code = apply_search_replace(streamlit_app_code, search_replace_pairs)
 
     return streamlit_app_code
 
-async def streamlit_handler(notebook_path: str, edit_prompt: str = "") -> Tuple[bool, Optional[str], str]:
+async def streamlit_handler(notebook_path: AbsoluteNotebookPath, edit_prompt: str = "") -> None:
     """Handler function for streamlit code generation and validation"""
 
-    clean_directory_check(notebook_path)
-
+    # Convert to absolute path for consistent handling
     notebook_code = parse_jupyter_notebook_to_extract_required_content(notebook_path)
-    app_directory = get_app_directory(notebook_path)
+    app_directory = get_absolute_notebook_dir_path(notebook_path)
+    app_path = get_absolute_app_path(app_directory)
     
     if edit_prompt != "":
         # If the user is editing an existing streamlit app, use the update function
-        streamlit_code = get_app_code_from_file(app_directory)
+        streamlit_code = get_app_code_from_file(app_path)
         
         if streamlit_code is None:
-            return False, '', "Error updating existing streamlit app because app.py file was not found."
+            raise StreamlitConversionError("Error updating existing streamlit app because app.py file was not found.", 404)
         
         streamlit_code = await update_existing_streamlit_code(notebook_code, streamlit_code, edit_prompt)
     else:
@@ -139,17 +133,13 @@ async def streamlit_handler(notebook_path: str, edit_prompt: str = "") -> Tuple[
         if has_validation_error:
             # TODO: We can't easily get the key type here, so for the beta release
             # we are just defaulting to the mito server key since that is by far the most common.
-            log_streamlit_app_creation_retry('mito_server_key', MessageType.STREAMLIT_CONVERSION, error)
+            log_streamlit_app_validation_retry('mito_server_key', MessageType.STREAMLIT_CONVERSION, errors)
         tries+=1
 
     if has_validation_error:
-        log_streamlit_app_creation_error('mito_server_key', MessageType.STREAMLIT_CONVERSION, error, edit_prompt)
-        return False, '', "Error generating streamlit code by agent"
+        final_errors = ', '.join(errors)
+        raise StreamlitConversionError(f"Streamlit agent failed generating code after max retries. Errors: {final_errors}", 500)
     
     # Finally, update the app.py file with the new code
-    success_flag, app_path, message = create_app_file(app_directory, streamlit_code)
-    if not success_flag:
-        log_streamlit_app_creation_error('mito_server_key', MessageType.STREAMLIT_CONVERSION, message, edit_prompt)
-    
-    log_streamlit_app_creation_success('mito_server_key', MessageType.STREAMLIT_CONVERSION, edit_prompt)
-    return success_flag, app_path, message
+    create_app_file(app_path, streamlit_code)
+    log_streamlit_app_conversion_success('mito_server_key', MessageType.STREAMLIT_CONVERSION, edit_prompt)

mito_ai/streamlit_conversion/streamlit_utils.py

@@ -3,9 +3,11 @@
 
 import re
 import json
-import os
 from typing import Dict, List, Optional, Tuple, Any
+from mito_ai.path_utils import AbsoluteAppPath, AbsoluteNotebookPath
 from pathlib import Path
+from mito_ai.utils.error_classes import StreamlitConversionError
+
 
 def extract_code_blocks(message_content: str) -> str:
     """
@@ -28,19 +30,7 @@ def extract_code_blocks(message_content: str) -> str:
     result = '\n'.join(matches)
     return result
 
-def extract_unified_diff_blocks(message_content: str) -> str:
-    """
-    Extract all unified_diff blocks from Claude's response.
-    """
-    if "```unified_diff" not in message_content:
-        return message_content
-    
-    pattern = r'```unified_diff\n(.*?)```'
-    matches = re.findall(pattern, message_content, re.DOTALL)
-    return '\n'.join(matches)
-
-
-def create_app_file(app_directory: str, code: str) -> Tuple[bool, str, str]:
+def create_app_file(app_path: AbsoluteAppPath, code: str) -> None:
     """
     Create app.py file and write code to it with error handling
 
@@ -53,40 +43,21 @@ def create_app_file(app_directory: str, code: str) -> Tuple[bool, str, str]:
 
     """
     try:
-        app_path = os.path.join(app_directory, "app.py")
-
         with open(app_path, 'w', encoding='utf-8') as f:
             f.write(code)
-        
-        return True, app_path, f"Successfully created {app_directory}"
     except IOError as e:
-        return False, '', f"Error creating file: {str(e)}"
-    except Exception as e:
-        return False, '', f"Unexpected error: {str(e)}"
+        raise StreamlitConversionError(f"Error creating app file: {str(e)}", 500)
     
-def get_app_code_from_file(app_directory: str) -> Optional[str]:
-    app_path = get_app_path(app_directory)
-    if app_path is None:
-        return None
+def get_app_code_from_file(app_path: AbsoluteAppPath) -> Optional[str]:
     with open(app_path, 'r', encoding='utf-8') as f:
         return f.read()
-    
-
-def get_app_path(app_directory: str) -> Optional[str]:
-    """
-    Check if the app.py file exists in the given directory.
-    """
-    app_path = os.path.join(app_directory, "app.py")
-    if not os.path.exists(app_path):
-        return None
-    return app_path
 
-def parse_jupyter_notebook_to_extract_required_content(notebook_path: str) -> List[Dict[str, Any]]:
+def parse_jupyter_notebook_to_extract_required_content(notebook_path: AbsoluteNotebookPath) -> List[Dict[str, Any]]:
     """
     Read a Jupyter notebook and filter cells to keep only cell_type and source fields.
 
     Args:
-        notebook_path (str): Path to the .ipynb file (can be relative or absolute)
+        notebook_path: Absolute path to the .ipynb file
 
     Returns:
         dict: Filtered notebook dictionary with only cell_type and source in cells
@@ -96,10 +67,6 @@ def parse_jupyter_notebook_to_extract_required_content(notebook_path: str) -> Li
         json.JSONDecodeError: If the file is not valid JSON
         KeyError: If the notebook doesn't have the expected structure
     """
-    # Convert to absolute path if it's not already absolute
-    # Handle both Unix-style absolute paths (starting with /) and Windows-style absolute paths
-    if not (notebook_path.startswith('/') or (len(notebook_path) > 1 and notebook_path[1] == ':')):
-        notebook_path = os.path.join(os.getcwd(), notebook_path)
     
     try:
         # Read the notebook file
@@ -108,7 +75,7 @@ def parse_jupyter_notebook_to_extract_required_content(notebook_path: str) -> Li
 
         # Check if 'cells' key exists
         if 'cells' not in notebook_data:
-            raise KeyError("Notebook does not contain 'cells' key")
+            raise StreamlitConversionError("Notebook does not contain 'cells' key", 400)
 
         # Filter each cell to keep only cell_type and source
         filtered_cells: List[Dict[str, Any]] = []
@@ -122,26 +89,6 @@ def parse_jupyter_notebook_to_extract_required_content(notebook_path: str) -> Li
         return filtered_cells
 
     except FileNotFoundError:
-        raise FileNotFoundError(f"Notebook file not found: {notebook_path}")
+        raise StreamlitConversionError(f"Notebook file not found: {notebook_path}", 404)
     except json.JSONDecodeError as e:
-        # JSONDecodeError requires msg, doc, pos
-        raise json.JSONDecodeError(f"Invalid JSON in notebook file: {str(e)}", e.doc if hasattr(e, 'doc') else '', e.pos if hasattr(e, 'pos') else 0)
-    except Exception as e:
-        raise Exception(f"Error processing notebook: {str(e)}")
-
-
-def resolve_notebook_path(notebook_path:str) -> str:
-    # Convert to absolute path if it's not already absolute
-    # Handle both Unix-style absolute paths (starting with /) and Windows-style absolute paths
-    if not (notebook_path.startswith('/') or (len(notebook_path) > 1 and notebook_path[1] == ':')):
-        notebook_path = os.path.join(os.getcwd(), notebook_path)
-    return notebook_path
-
-def clean_directory_check(notebook_path: str) -> None:
-    notebook_path = resolve_notebook_path(notebook_path)
-    # pathlib handles the cross OS path conversion automatically
-    path = Path(notebook_path).resolve()
-    dir_path = path.parent
-
-    if not dir_path.exists():
-        raise ValueError(f"Directory does not exist: {dir_path}")
+        raise StreamlitConversionError(f"Invalid JSON in notebook file: {str(e)}", 400)

mito_ai/streamlit_conversion/validate_streamlit_app.py

@@ -1,28 +1,18 @@
 # Copyright (c) Saga Inc.
 # Distributed under the terms of the GNU Affero General Public License v3.0 License.
 
-import sys
 import os
-import time
-import requests
 import tempfile
-import shutil
 import traceback
 import ast
-import importlib.util
 import warnings
 from typing import List, Tuple, Optional, Dict, Any, Generator
 from streamlit.testing.v1 import AppTest
 from contextlib import contextmanager
-from mito_ai.streamlit_conversion.streamlit_utils import resolve_notebook_path
-
-
-# warnings.filterwarnings("ignore", message=r".*missing ScriptRunContext.*")
-# warnings.filterwarnings("ignore", category=UserWarning)
+from mito_ai.path_utils import AbsoluteNotebookPath, get_absolute_notebook_dir_path
 
 warnings.filterwarnings("ignore", message=".*bare mode.*")
 
-
 def get_syntax_error(app_code: str) -> Optional[str]:
     """Check if the Python code has valid syntax"""
     try:
@@ -32,10 +22,10 @@ def get_syntax_error(app_code: str) -> Optional[str]:
         error_msg = ''.join(traceback.format_exception(type(e), e, e.__traceback__))
         return error_msg
 
-def get_runtime_errors(app_code: str, app_path: str) -> Optional[List[Dict[str, Any]]]:
+def get_runtime_errors(app_code: str, app_path: AbsoluteNotebookPath) -> Optional[List[Dict[str, Any]]]:
     """Start the Streamlit app in a subprocess"""  
     
-    directory = os.path.dirname(app_path)
+    directory = get_absolute_notebook_dir_path(app_path)
     
     @contextmanager
     def change_working_directory(path: str) -> Generator[None, Any, None]:
@@ -89,7 +79,7 @@ def get_runtime_errors(app_code: str, app_path: str) -> Optional[List[Dict[str,
             except OSError:
                 pass  # File might already be deleted
 
-def check_for_errors(app_code: str, app_path: str) -> List[Dict[str, Any]]:
+def check_for_errors(app_code: str, app_path: AbsoluteNotebookPath) -> List[Dict[str, Any]]:
     """Complete validation pipeline"""
     errors: List[Dict[str, Any]] = []
 
@@ -100,7 +90,6 @@ def check_for_errors(app_code: str, app_path: str) -> List[Dict[str, Any]]:
             errors.append({'type': 'syntax', 'details': syntax_error})
 
         runtime_errors = get_runtime_errors(app_code, app_path)
-        
         if runtime_errors:
             errors.extend(runtime_errors)
         
@@ -109,10 +98,8 @@ def check_for_errors(app_code: str, app_path: str) -> List[Dict[str, Any]]:
 
     return errors
 
-def validate_app(app_code: str, notebook_path: str) -> Tuple[bool, List[str]]:
+def validate_app(app_code: str, notebook_path: AbsoluteNotebookPath) -> Tuple[bool, List[str]]:
     """Convenience function to validate Streamlit code"""
-    notebook_path = resolve_notebook_path(notebook_path)
-    
     errors = check_for_errors(app_code, notebook_path)
     
     has_validation_error = len(errors) > 0

mito_ai/streamlit_preview/handlers.py

@@ -2,17 +2,17 @@
 # Distributed under the terms of the GNU Affero General Public License v3.0 License.
 
 import os
-import tempfile
+from typing import Literal, TypedDict
 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
 import tornado
 from jupyter_server.base.handlers import APIHandler
-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
-
+from mito_ai.path_utils import  AbsoluteNotebookPath, get_absolute_notebook_dir_path, get_absolute_notebook_path
+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
+import traceback
 
 
 class StreamlitPreviewHandler(APIHandler):
@@ -22,102 +22,62 @@ class StreamlitPreviewHandler(APIHandler):
         """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)
-
     @tornado.web.authenticated
     async def post(self) -> None:
         """Start a new streamlit preview."""
         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, force_recreate, edit_prompt = validate_request_body(body)
 
             # Ensure app exists
-            resolved_notebook_path = self._resolve_notebook_path(notebook_path)
-
-            success, error_msg = await ensure_app_exists(resolved_notebook_path, force_recreate, edit_prompt)
-
-            if not success:
-                self.set_status(500)
-                self.finish({"error": error_msg})
-                return
+            absolute_notebook_path = get_absolute_notebook_path(notebook_path)
+            await ensure_app_exists(absolute_notebook_path, force_recreate, 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
+            absolute_app_directory = get_absolute_notebook_dir_path(absolute_notebook_path)
+            port = self.preview_manager.start_streamlit_preview(absolute_app_directory, 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

@@ -11,6 +11,7 @@ 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 +38,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 +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}")
+                print(f"Waiting for app to be ready...")
                 pass
             
             time.sleep(1)
@@ -153,7 +154,6 @@ class StreamlitPreviewManager:
 # Global instance
 _preview_manager = StreamlitPreviewManager()
 
-
 def get_preview_manager() -> StreamlitPreviewManager:
     """Get the global preview manager instance."""
     return _preview_manager