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

mito_ai/streamlit_conversion/prompts/streamlit_error_correction_prompt.py

@@ -1,7 +1,7 @@
 # Copyright (c) Saga Inc.
 # Distributed under the terms of the GNU Affero General Public License v3.0 License.
 
-from mito_ai.streamlit_conversion.prompts.prompt_constants import unified_diff_instructions
+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_streamlit_error_correction_prompt(error: str, streamlit_app_code: str) -> str:
@@ -12,7 +12,7 @@ def get_streamlit_error_correction_prompt(error: str, streamlit_app_code: str) -
 
 Your job is to fix the error now. Only fix the specific error that you are instructed to fix now. Do not fix other error that that you anticipate. You will be asked to fix other errors later.
 
-{unified_diff_instructions}
+{search_replace_instructions}
 
 ===============================================
 
@@ -21,7 +21,7 @@ EXISTING STREAMLIT APP:
 
 ===============================================
 
-Please create a unified diff that corrects this error. Please keep your fix concise:
+Please create a search/replace block that corrects this error. Please keep your fix concise:
 {error}
 
 """

mito_ai/streamlit_conversion/prompts/streamlit_finish_todo_prompt.py

@@ -2,7 +2,7 @@
 # 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 MITO_TODO_PLACEHOLDER, unified_diff_instructions
+from mito_ai.streamlit_conversion.prompts.prompt_constants import MITO_TODO_PLACEHOLDER, search_replace_instructions
 from mito_ai.streamlit_conversion.prompts.prompt_utils import add_line_numbers_to_code
 
 def get_finish_todo_prompt(notebook: List[dict], existing_streamlit_app_code: str, todo_placeholder: str) -> str:
@@ -25,7 +25,7 @@ You have ONE and ONLY ONE opportunity to complete this TODO. If you do not finis
 - 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.
 
-{unified_diff_instructions}
+{search_replace_instructions}
 
 ===============================================
 

mito_ai/streamlit_conversion/prompts/update_existing_app_prompt.py

@@ -2,7 +2,7 @@
 # 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 unified_diff_instructions
+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:
@@ -30,7 +30,7 @@ You have ONE and ONLY ONE opportunity to complete this edit request. If you do n
 - 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.
 
-{unified_diff_instructions}
+{search_replace_instructions}
 
 ===============================================
 

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

@@ -1,27 +1,20 @@
 # Copyright (c) Saga Inc.
 # Distributed under the terms of the GNU Affero General Public License v3.0 License.
 
-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 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, 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, 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"""
@@ -57,10 +50,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 +72,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 +94,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, app_file_name: AppFileName, 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, 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_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:
@@ -128,28 +121,24 @@ async def streamlit_handler(notebook_path: str, edit_prompt: str = "") -> Tuple[
         streamlit_code = await generate_new_streamlit_code(notebook_code)
        
     # Then, after creating/updating the app, validate that the new code runs 
-    has_validation_error, errors = validate_app(streamlit_code, notebook_path)
+    errors = validate_app(streamlit_code, notebook_path)
     tries = 0
-    while has_validation_error and tries < 5:
+    while len(errors)>0 and tries < 5:
         for error in errors:
             streamlit_code = await correct_error_in_generation(error, streamlit_code)
         
-        has_validation_error, errors = validate_app(streamlit_code, notebook_path)
+        errors = validate_app(streamlit_code, notebook_path)
         
-        if has_validation_error:
+        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_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"
+    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
-    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,10 @@
 
 import re
 import json
-import os
-from typing import Dict, List, Optional, Tuple, Any
-from pathlib import Path
+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:
     """
@@ -28,65 +29,26 @@ 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.
+def create_app_file(app_path: AbsoluteAppPath, code: str) -> None:
     """
-    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]:
-    """
-    Create app.py file and write code to it with error handling
-
-    Args:
-        file_path (str): The actual content from the agent's response
-        code (str): The actual content from the agent's response
-
-    Returns:
-        str: Removes the ```python``` part to be able to parse the code
-
+    Create .py file and write code to it with error handling
     """
     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 +58,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 +66,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 +80,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,12 +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) -> 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
     stringified_errors = [str(error) for error in errors]
-    return has_validation_error, stringified_errors
+    return stringified_errors

mito_ai/streamlit_preview/__init__.py

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