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

mito_ai/streamlit_conversion/prompts/update_existing_app_prompt.py

@@ -0,0 +1,50 @@
+# Copyright (c) Saga Inc.
+# Distributed under the terms of the GNU Affero General Public License v3.0 License.
+
+from typing import List
+from mito_ai.streamlit_conversion.prompts.prompt_constants import search_replace_instructions
+from mito_ai.streamlit_conversion.prompts.prompt_utils import add_line_numbers_to_code
+
+def get_update_existing_app_prompt(notebook: List[dict], streamlit_app_code: str, edit_prompt: str) -> str:
+    """
+    This prompt is used to update an existing streamlit app.
+    """
+    
+    existing_streamlit_app_code_with_line_numbers = add_line_numbers_to_code(streamlit_app_code)
+    
+    return f"""
+
+GOAL: You've previously created a first draft of the Streamlit app. Now the user reviewed it and provided feedback.Update the existing streamlit app according to the feedback provided by the user. Use the input notebook to help you understand what code needs to be added, changed, or modified to fulfill the user's edit request.
+
+**CRITICAL COMPLETION REQUIREMENT:**
+You have ONE and ONLY ONE opportunity to complete this edit request. If you do not finish the entire task completely, the application will be broken and unusable. This is your final chance to get it right.
+
+**COMPLETION RULES:**
+1. **NEVER leave partial work** - If the edit request requires generating a list with 100 items, provide ALL 100 items.
+2. **NEVER use placeholders** - This is your only opportunity to fulfill this edit request, so do not leave yourself another TODOs.
+3. **NEVER assume "good enough"** - Complete the task to 100% satisfaction.
+4. **If the task seems large, that's exactly why it needs to be done now** - This is your only chance
+
+**HOW TO DETERMINE IF TASK IS COMPLETE:**
+- If building a list/dictionary: Include ALL items that should be in the final data structure.
+- If creating functions: Implement ALL required functionality.
+- If converting a visualization: Copy over ALL of the visualization code from the notebook, including all styling and formatting.
+
+{search_replace_instructions}
+
+===============================================
+
+INPUT NOTEBOOK: 
+{notebook}
+
+===============================================
+
+EXISTING STREAMLIT APP:
+{existing_streamlit_app_code_with_line_numbers}
+
+===============================================
+
+USER EDIT REQUEST:
+{edit_prompt}
+
+"""

mito_ai/streamlit_conversion/search_replace_utils.py

@@ -0,0 +1,94 @@
+# 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
+from mito_ai.utils.telemetry_utils import log
+
+
+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))
+            log("mito_ai_search_text_found_multiple_times", params={"search_text": repr(search_text)}, key_type="mito_server_key")
+
+        # Perform the replacement
+        result = result.replace(search_text, replace_text, 1)
+    
+    return result

mito_ai/streamlit_conversion/streamlit_agent_handler.py

@@ -0,0 +1,144 @@
+# Copyright (c) Saga Inc.
+# Distributed under the terms of the GNU Affero General Public License v3.0 License.
+
+from anthropic.types import MessageParam
+from typing import List, cast
+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, 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.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, AppFileName, get_absolute_notebook_dir_path, get_absolute_app_path, get_app_file_name
+
+async def generate_new_streamlit_code(notebook: List[dict]) -> str:
+    """Send a query to the agent, get its response and parse the code"""
+    
+    prompt_text = get_streamlit_app_creation_prompt(notebook)
+    
+    messages: List[MessageParam] = [
+        cast(MessageParam, {
+            "role": "user",
+            "content": [{
+                "type": "text",
+                "text": prompt_text
+            }]
+        })
+    ]
+    agent_response = await get_response_from_agent(messages)
+    converted_code = extract_code_blocks(agent_response)
+    
+    # Extract the TODOs from the agent's response
+    todo_placeholders = extract_todo_placeholders(agent_response)
+    
+    for todo_placeholder in todo_placeholders:
+        print(f"Processing AI TODO: {todo_placeholder}")
+        todo_prompt = get_finish_todo_prompt(notebook, converted_code, todo_placeholder)
+        todo_messages: List[MessageParam] = [
+            cast(MessageParam, {
+                "role": "user",
+                "content": [{
+                    "type": "text",
+                    "text": todo_prompt
+                }]
+            })
+        ]
+        todo_response = await get_response_from_agent(todo_messages)
+        
+        # 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
+
+
+async def update_existing_streamlit_code(notebook: List[dict], streamlit_app_code: str, edit_prompt: str) -> str:
+    """Send a query to the agent, get its response and parse the code"""
+    prompt_text = get_update_existing_app_prompt(notebook, streamlit_app_code, edit_prompt)
+    
+    messages: List[MessageParam] = [
+        cast(MessageParam, {
+            "role": "user",
+            "content": [{
+                "type": "text",
+                "text": prompt_text
+            }]
+        })
+    ]
+    
+    agent_response = await get_response_from_agent(messages)
+    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
+
+
+async def correct_error_in_generation(error: str, streamlit_app_code: str) -> str:
+    """If errors are present, send it back to the agent to get corrections in code"""
+    messages: List[MessageParam] = [
+        cast(MessageParam, {
+            "role": "user",
+            "content": [{
+                "type": "text",
+                "text": get_streamlit_error_correction_prompt(error, streamlit_app_code)
+            }]
+        })
+    ]
+    agent_response = await get_response_from_agent(messages)
+    
+    # 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: AbsoluteNotebookPath, app_file_name: AppFileName, edit_prompt: str = "") -> None:
+    """Handler function for streamlit code generation and validation"""
+
+    # Convert to absolute path for consistent handling
+    notebook_code = parse_jupyter_notebook_to_extract_required_content(notebook_path)
+    app_directory = get_absolute_notebook_dir_path(notebook_path)
+    app_path = get_absolute_app_path(app_directory, app_file_name)
+    
+    if edit_prompt != "":
+        # If the user is editing an existing streamlit app, use the update function
+        streamlit_code = get_app_code_from_file(app_path)
+        
+        if streamlit_code is None:
+            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:
+        # Otherwise generate a new streamlit app
+        streamlit_code = await generate_new_streamlit_code(notebook_code)
+       
+    # Then, after creating/updating the app, validate that the new code runs 
+    errors = validate_app(streamlit_code, notebook_path)
+    tries = 0
+    while len(errors)>0 and tries < 5:
+        for error in errors:
+            streamlit_code = await correct_error_in_generation(error, streamlit_code)
+        
+        errors = validate_app(streamlit_code, notebook_path)
+        
+        if len(errors)>0:
+            # 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_validation_retry('mito_server_key', MessageType.STREAMLIT_CONVERSION, errors)
+        tries+=1
+
+    if len(errors)>0:
+        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
+    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

@@ -0,0 +1,85 @@
+# Copyright (c) Saga Inc.
+# Distributed under the terms of the GNU Affero General Public License v3.0 License.
+
+import re
+import json
+from typing import Dict, List, Optional, Any
+from mito_ai.path_utils import AbsoluteAppPath, AbsoluteNotebookPath
+from mito_ai.utils.error_classes import StreamlitConversionError
+
+
+def extract_code_blocks(message_content: str) -> str:
+    """
+    Extract all code blocks from Claude's response.
+
+    Args:
+        message_content (str): The actual content from the agent's response
+
+    Returns:
+        str: Removes the ```python``` part to be able to parse the code
+    """
+    if "```python" not in message_content:
+        return message_content
+
+    # Use regex to find all Python code blocks
+    pattern = r'```python\n(.*?)```'
+    matches = re.findall(pattern, message_content, re.DOTALL)
+
+    # Concatenate with single newlines
+    result = '\n'.join(matches)
+    return result
+
+def create_app_file(app_path: AbsoluteAppPath, code: str) -> None:
+    """
+    Create .py file and write code to it with error handling
+    """
+    try:
+        with open(app_path, 'w', encoding='utf-8') as f:
+            f.write(code)
+    except IOError as e:
+        raise StreamlitConversionError(f"Error creating app file: {str(e)}", 500)
+    
+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 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: Absolute path to the .ipynb file
+
+    Returns:
+        dict: Filtered notebook dictionary with only cell_type and source in cells
+
+    Raises:
+        FileNotFoundError: If the notebook file doesn't exist
+        json.JSONDecodeError: If the file is not valid JSON
+        KeyError: If the notebook doesn't have the expected structure
+    """
+    
+    try:
+        # Read the notebook file
+        with open(notebook_path, 'r', encoding='utf-8') as f:
+            notebook_data: Dict[str, Any] = json.load(f)
+
+        # Check if 'cells' key exists
+        if 'cells' not in notebook_data:
+            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]] = []
+        for cell in notebook_data['cells']:
+            filtered_cell: Dict[str, Any] = {
+                'cell_type': cell.get('cell_type', ''),
+                'source': cell.get('source', [])
+            }
+            filtered_cells.append(filtered_cell)
+            
+        return filtered_cells
+
+    except FileNotFoundError:
+        raise StreamlitConversionError(f"Notebook file not found: {notebook_path}", 404)
+    except json.JSONDecodeError as e:
+        raise StreamlitConversionError(f"Invalid JSON in notebook file: {str(e)}", 400)

mito_ai/streamlit_conversion/validate_streamlit_app.py

@@ -0,0 +1,105 @@
+# Copyright (c) Saga Inc.
+# Distributed under the terms of the GNU Affero General Public License v3.0 License.
+
+import os
+import tempfile
+import traceback
+import ast
+import warnings
+from typing import List, Tuple, Optional, Dict, Any, Generator
+from streamlit.testing.v1 import AppTest
+from contextlib import contextmanager
+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:
+        ast.parse(app_code)
+        return None
+    except SyntaxError as e:
+        error_msg = ''.join(traceback.format_exception(type(e), e, e.__traceback__))
+        return error_msg
+
+def get_runtime_errors(app_code: str, app_path: AbsoluteNotebookPath) -> Optional[List[Dict[str, Any]]]:
+    """Start the Streamlit app in a subprocess"""  
+    
+    directory = get_absolute_notebook_dir_path(app_path)
+    
+    @contextmanager
+    def change_working_directory(path: str) -> Generator[None, Any, None]:
+        """
+        Context manager to temporarily change working directory
+        so that relative paths are still valid when we run the app
+        """
+        if path == '':
+            yield
+        
+        original_cwd = os.getcwd()
+        try:
+            os.chdir(path)
+            yield
+        finally:
+            os.chdir(original_cwd)
+    
+    with change_working_directory(directory):
+        # Create a temporary file that uses UTF-8 encoding so 
+        # we don't run into issues with non-ASCII characters on Windows.
+        # We use utf-8 encoding when writing the app.py file so this validation
+        # code mirrors the actual file. 
+
+        # Note: Since the AppTest.from_file tries to open the file, we need to first close the file
+        # by exiting the context manager and using the delete=False flag so that the file still exists.
+        # Windows can't open the same file twice at the same time. We cleanup at the end.
+        with tempfile.NamedTemporaryFile(mode="w", suffix=".py", delete=False, encoding="utf-8") as f:
+            f.write(app_code)
+            temp_path = f.name
+
+        try:
+            # Run Streamlit test from file with UTF-8 encoding
+            app_test = AppTest.from_file(temp_path, default_timeout=30)
+            app_test.run()
+            
+            # Check for exceptions
+            if app_test.exception:
+                errors = [{'type': 'exception', 'details': exc.value, 'message': exc.message, 'stack_trace': exc.stack_trace} for exc in app_test.exception]
+                return errors
+                    
+            # Check for error messages
+            if app_test.error:
+                errors = [{'type': 'error', 'details': err.value} for err in app_test.error]
+                return errors
+            
+            return None
+        finally:
+            # Clean up the temporary file
+            try:
+                os.unlink(temp_path)
+            except OSError:
+                pass  # File might already be deleted
+
+def check_for_errors(app_code: str, app_path: AbsoluteNotebookPath) -> List[Dict[str, Any]]:
+    """Complete validation pipeline"""
+    errors: List[Dict[str, Any]] = []
+
+    try:
+        # Step 1: Check syntax
+        syntax_error = get_syntax_error(app_code)
+        if syntax_error:
+            errors.append({'type': 'syntax', 'details': syntax_error})
+
+        runtime_errors = get_runtime_errors(app_code, app_path)
+        if runtime_errors:
+            errors.extend(runtime_errors)
+        
+    except Exception as e:
+        errors.append({'type': 'validation', 'details': str(e)})
+
+    return errors
+
+def validate_app(app_code: str, notebook_path: AbsoluteNotebookPath) -> List[str]:
+    """Convenience function to validate Streamlit code"""
+    errors = check_for_errors(app_code, notebook_path)
+    stringified_errors = [str(error) for error in errors]
+    return stringified_errors

mito_ai/streamlit_preview/__init__.py

@@ -0,0 +1,6 @@
+# Copyright (c) Saga Inc.
+# Distributed under the terms of the GNU Affero General Public License v3.0 License.
+
+from .handlers import StreamlitPreviewHandler
+
+__all__ = ['StreamlitPreviewHandler']

mito_ai/streamlit_preview/handlers.py

@@ -0,0 +1,111 @@
+# Copyright (c) Saga Inc.
+# Distributed under the terms of the GNU Affero General Public License v3.0 License.
+
+import uuid
+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
+import traceback
+
+
+class StreamlitPreviewHandler(APIHandler):
+    """REST handler for streamlit preview operations."""
+
+    def initialize(self) -> None:
+        """Initialize the handler."""
+        self.preview_manager = StreamlitPreviewManager()
+
+    @tornado.web.authenticated
+    async def post(self) -> None:
+        """Start a new streamlit preview."""
+        try:
+            # Parse and validate request
+            body = self.get_json_body()
+            notebook_path, notebook_id, force_recreate, edit_prompt = validate_request_body(body)
+
+            # Ensure app exists
+            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)
+
+            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")
+
+                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())
+            port = self.preview_manager.start_streamlit_preview(absolute_notebook_dir_path, app_file_name, preview_id)
+
+            # Return success response
+            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"Exception in streamlit preview handler: {e}")
+            self.set_status(500)
+            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:
+        """Stop a streamlit preview."""
+        try:
+            if not preview_id:
+                self.set_status(400)
+                self.finish({"error": "Missing preview_id parameter"})
+                return
+
+            # Stop the preview
+            stopped = self.preview_manager.stop_preview(preview_id)
+
+            if stopped:
+                self.set_status(204)  # No content
+            else:
+                self.set_status(404)
+                self.finish({"error": f"Preview {preview_id} not found"})
+
+        except Exception as e:
+            self.set_status(500)
+            self.finish({"error": f"Internal server error: {str(e)}"})