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

mito_ai/completions/prompt_builders/prompt_constants.py

@@ -19,6 +19,7 @@ ACTIVE_CELL_ID_SECTION_HEADING = "The ID of the active code cell:"
 ACTIVE_CELL_OUTPUT_SECTION_HEADING = "Output of the active code cell:"
 GET_CELL_OUTPUT_TOOL_RESPONSE_SECTION_HEADING = "Output of the code cell you just applied the CELL_UPDATE to:"
 JUPYTER_NOTEBOOK_SECTION_HEADING = "Jupyter Notebook:"
+STREAMLIT_APP_STATUS_SECTION_HEADING = "Streamlit App Status:"
 
 # Placeholder text used when trimming content from messages
 CONTENT_REMOVED_PLACEHOLDER = "Content removed to save space" 

mito_ai/completions/prompt_builders/utils.py

@@ -69,3 +69,16 @@ def get_selected_context_str(additional_context: Optional[List[Dict[str, str]]])
 
     # STEP 3: Combine into a single string
     return "\n\n".join(context_parts)
+
+
+def get_streamlit_app_status_str(notebook_id: str, notebook_path: str) -> str:
+    """
+    Get the streamlit app status string.
+    """
+    from mito_ai.path_utils import does_notebook_id_have_corresponding_app
+    if does_notebook_id_have_corresponding_app(notebook_id, notebook_path):
+        return "The notebook has an existing Streamlit app that you can edit"
+    return "The notebook does not have an existing Streamlit app. If you want to show an app to the user, you must create a new one."
+
+
+

mito_ai/path_utils.py

@@ -0,0 +1,70 @@
+# Copyright (c) Saga Inc.
+# Distributed under the terms of the GNU Affero General Public License v3.0 License.
+
+from typing import NewType
+import os
+from mito_ai.utils.error_classes import StreamlitPreviewError
+
+# Type definitions for better type safety
+AbsoluteNotebookPath = NewType('AbsoluteNotebookPath', str)
+AbsoluteNotebookDirPath = NewType('AbsoluteNotebookDirPath', str)
+AbsoluteAppPath = NewType('AbsoluteAppPath', str)
+AppFileName = NewType("AppFileName", str)
+
+def get_absolute_notebook_path(notebook_path: str) -> AbsoluteNotebookPath:
+    """
+    Convert any notebook path to an absolute path.
+    
+    Args:
+        notebook_path: Path to the notebook (can be relative or absolute)
+        
+    Returns:
+        AbsoluteNotebookPath: The absolute path to the notebook
+        
+    Raises:
+        ValueError: If the path is invalid or empty
+    """
+    if not notebook_path or not notebook_path.strip():
+        raise StreamlitPreviewError("Notebook path cannot be empty", 400)
+    
+    absolute_path = os.path.abspath(notebook_path)
+    return AbsoluteNotebookPath(absolute_path)
+    
+
+def get_absolute_notebook_dir_path(notebook_path: AbsoluteNotebookPath) -> AbsoluteNotebookDirPath:
+    """
+    Get the absolute directory containing the notebook.
+    """
+    return AbsoluteNotebookDirPath(os.path.dirname(notebook_path))
+
+def get_absolute_app_path(app_directory: AbsoluteNotebookDirPath, app_file_name: AppFileName) -> AbsoluteAppPath:
+    """
+    Get the absolute path to the app
+    """
+    return AbsoluteAppPath(os.path.join(app_directory, app_file_name))
+
+def get_app_file_name(notebook_id: str) -> AppFileName:
+    """
+    Converts the notebook id into the corresponding app id
+    """
+    mito_app_name = notebook_id.replace('mito-notebook-', 'mito-app-')
+    return AppFileName(f'{mito_app_name}.py')
+
+def does_app_path_exist(app_path: AbsoluteAppPath) -> bool:
+    """
+    Check if the app file exists
+    """
+    return os.path.exists(app_path)
+
+def does_notebook_id_have_corresponding_app(notebook_id: str, notebook_path: str) -> bool:
+    """
+    Given a notebook_id and raw notebook_path checks if the notebook has a corresponding
+    app by converting the notebook_path into an absolute path and converting the notebook_id
+    into an app name
+    """
+
+    app_file_name = get_app_file_name(notebook_id)
+    notebook_path = get_absolute_notebook_path(notebook_path)
+    app_directory = get_absolute_notebook_dir_path(notebook_path)
+    app_path = get_absolute_app_path(app_directory, app_file_name)
+    return does_app_path_exist(app_path)

mito_ai/streamlit_conversion/agent_utils.py

@@ -1,221 +1,24 @@
 # Copyright (c) Saga Inc.
 # Distributed under the terms of the GNU Affero General Public License v3.0 License.
 
-from typing import List
+from typing import List, Tuple
 import re
 from anthropic.types import MessageParam
-from mito_ai.streamlit_conversion.streamlit_system_prompt import streamlit_system_prompt
+from mito_ai.streamlit_conversion.prompts.streamlit_system_prompt import streamlit_system_prompt
 from mito_ai.utils.anthropic_utils import stream_anthropic_completion_from_mito_server
-from unidiff import PatchSet
 from mito_ai.streamlit_conversion.prompts.prompt_constants import MITO_TODO_PLACEHOLDER
 from mito_ai.completions.models import MessageType
 
-STREAMLIT_AI_MODEL = "claude-3-5-haiku-latest"
+STREAMLIT_AI_MODEL = "claude-sonnet-4-5-20250929"
 
 def extract_todo_placeholders(agent_response: str) -> List[str]:
     """Extract TODO placeholders from the agent's response"""
     return [line.strip() for line in agent_response.split('\n') if MITO_TODO_PLACEHOLDER in line]
 
-
-def apply_patch_to_text(text: str, diff: str) -> str:
-    """
-    Apply a *unified-diff* (git-style) patch to the given text and return
-    the updated contents.
-
-    Parameters
-    ----------
-    text : str
-        The original file contents.
-    diff : str
-        A unified diff that transforms *text* into the desired output.
-        The diff must reference exactly one file (the Streamlit app).
-        NOTE: This assumes a custom format where BOTH -X,Y and +X,Y
-        reference the original file line numbers.
-
-    Returns
-    -------
-    str
-        The patched contents.
-
-    Raises
-    ------
-    ValueError
-        If the patch cannot be applied or references more than one file.
-    """
-    # Nothing to do
-    if not diff.strip():
-        return text
-
-    # Parse the patch
-    patch = PatchSet(diff.splitlines(keepends=True))
-
-    # We expect a single-file patch (what the prompt asks the model to emit)
-    if len(patch) == 0:
-        raise ValueError("No patches found in diff")
-    
-    # Check that all patches are for the same file
-    file_names = set(p.source_file for p in patch)
-    if len(file_names) > 1:
-        raise ValueError(
-            f"Expected patches for exactly one file, got files: {file_names}"
-        )
-
-    # Apply all hunks from all patches (they should all be for the same file)
-    original_lines = text.splitlines(keepends=True)
-    result_lines: List[str] = []
-    cursor = 0  # index in original_lines (0-based)
-
-    # Process all hunks from all patches
-    # We only expect one patch file, but it always returns as a list 
-    # so we just iterate over it
-    for file_patch in patch:
-        for hunk in file_patch:
-            # Since hunks reference the original file, just convert to 0-based
-            hunk_start = hunk.source_start - 1
-            
-            # Copy unchanged lines before this hunk
-            while cursor < hunk_start:
-                if cursor < len(original_lines):
-                    result_lines.append(original_lines[cursor])
-                cursor += 1
-
-            # Apply hunk line-by-line
-            for line in hunk:
-                if line.is_context:
-                    # Use the line from the diff to preserve exact formatting
-                    result_lines.append(line.value)
-                    cursor += 1
-                elif line.is_removed:
-                    cursor += 1  # Skip this line from the original
-                elif line.is_added:
-                    # Use the line from the diff to preserve exact formatting
-                    result_lines.append(line.value)
-
-    # Copy any remaining lines after the last hunk
-    result_lines.extend(original_lines[cursor:])
-
-    return "".join(result_lines)
-
-
-def fix_context_lines(diff: str) -> str:
-    """
-    Fix context lines in unified diff to ensure they all start with a space character.
-    
-    In unified diffs, context lines (unchanged lines) must start with a single space ' ',
-    even if the line itself is empty. The AI sometimes generates diffs where empty
-    context lines are just blank lines without the leading space, which causes the
-    unidiff parser to fail.
-    
-    Args:
-        diff (str): The unified diff string
-        
-    Returns:
-        str: The corrected diff with proper context line formatting
-    """
-    lines = diff.split('\n')
-    corrected_lines = []
-    in_hunk = False
-    
-    for i, line in enumerate(lines):
-        # Check if we're entering a hunk
-        if line.startswith('@@'):
-            in_hunk = True
-            corrected_lines.append(line)
-            continue
-            
-        # Check if we're leaving a hunk (new file header)
-        if line.startswith('---') or line.startswith('+++'):
-            in_hunk = False
-            corrected_lines.append(line)
-            continue
-            
-        if in_hunk:
-            # We're inside a hunk
-            if line.startswith(' ') or line.startswith('-') or line.startswith('+'):
-                # Already has proper diff marker
-                corrected_lines.append(line)
-            elif line.strip() == '':
-                # Empty line should be a context line with leading space
-                corrected_lines.append(' ')
-            else:
-                # Line without diff marker - treat as context line
-                corrected_lines.append(' ' + line)
-        else:
-            # Outside hunk - keep as is
-            corrected_lines.append(line)
-    
-    return '\n'.join(corrected_lines)
-    
-
-def fix_diff_headers(diff: str) -> str:
-    """
-    The AI is generally not very good at counting the number of lines in the diff. If the hunk header has
-    an incorrect count, then the patch will fail. So instead we just calculate the counts ourselves, its deterministic.
-    
-    If no header is provided at all, then there is nothing to fix.
-    """
-    # First fix context lines to ensure they have proper leading spaces
-    diff = fix_context_lines(diff)
-    
-    lines = diff.split('\n')
-    
-    for i, line in enumerate(lines):
-        if line.startswith('@@'):
-            # Extract the starting line numbers
-            match = re.match(r'@@ -(\d+),\d+ \+(\d+),\d+ @@', line)
-            if match:
-                old_start = match.group(1)
-                new_start = match.group(2)
-                
-                # Count lines in this hunk
-                old_count = 0
-                new_count = 0
-                
-                # Find the end of this hunk (next @@ line or end of file)
-                hunk_end = len(lines)
-                for j in range(i + 1, len(lines)):
-                    if lines[j].startswith('@@'):
-                        hunk_end = j
-                        break
-                
-                # Count lines in this hunk
-                for j in range(i + 1, hunk_end):
-                    hunk_line = lines[j]
-                    # Empty lines are treated as context lines
-                    if hunk_line == '' or hunk_line.startswith(' ') or hunk_line.startswith('-'):
-                        old_count += 1
-                    if hunk_line == '' or hunk_line.startswith(' ') or hunk_line.startswith('+'):
-                        new_count += 1
-                
-                # Replace the header with correct counts
-                lines[i] = f"@@ -{old_start},{old_count} +{new_start},{new_count} @@"
-    
-    corrected_diff = '\n'.join(lines)
-    corrected_diff = corrected_diff.lstrip()
-    
-    # If there is no diff, just return it without fixing file headers
-    if len(corrected_diff) == 0:
-        return corrected_diff
-    
-    # Remove known problametic file component headers that the AI sometimes returns
-    problamatic_file_header_components = ['--- a/app.py +++ b/app.py']
-    for problamatic_file_header_component in problamatic_file_header_components:
-        corrected_diff = corrected_diff.removeprefix(problamatic_file_header_component).lstrip()
-        
-    
-    # If the diff is missing the file component of the header, add it
-    valid_header_component = """--- a/app.py
-+++ b/app.py"""
-    if not corrected_diff.startswith(valid_header_component):
-        corrected_diff = valid_header_component + '\n' + corrected_diff
-    
-    return corrected_diff
-
-
 async def get_response_from_agent(message_to_agent: List[MessageParam]) -> str:
     """Gets the streaming response from the agent using the mito server"""
     model = STREAMLIT_AI_MODEL
-    max_tokens = 8192 # 64_000
+    max_tokens = 64000 # TODO: If we move to haiku, we must reset this to 8192
     temperature = 0.2
 
     accumulated_response = ""

mito_ai/streamlit_conversion/prompts/prompt_constants.py

@@ -3,121 +3,154 @@
 
 MITO_TODO_PLACEHOLDER = "# MITO_TODO_PLACEHOLDER"
 
-unified_diff_instructions = f"""
-RESPONSE FORMAT: Return the changes you want to make to the streamlit app as a **unified diff (git-style patch)**:
-
-A unified diff looks is the following and tells the system which lines of code to add, remove, or modify:
---- a/app.py
-+++ b/app.py
-@@ -START_LINE,1 +START_LINE,1 @@
- x = 1
- -y = 2
- +y = 3
- 
-The components of the unified diff are the following:
-- `--- a/app.py` -> The original file. We will always use the file app.py
-- `+++ b/app.py` -> The modified file. We will always use the file app.py
-- `@@ -START_LINE,1 +START_LINE,1 @@` -> The hunk header
-- `x = 1` -> The original context line
-- `-y = 2` -> The removed line
-- `+y = 3` -> The added line
-
-When you create a unified diff, you must follow the following format:
-- Begin with a ```unified_diff marker and a ``` end marker.
-- Always, include the standard header. On line 1: `--- a/app.py` and on line 2: `+++ b/app.py` like in the example above.
-- Show only the modified hunks; each hunk must start with an `@@` header with line numbers.
-- Within each hunk:
-  * Unchanged context lines start with a single space ` `.
-  * Removed lines start with `-`.
-  * Added lines start with `+`.
-- If there are **no changes**, return an empty string.
-- Do not include the line numbers in your response.
-
-**CRITICAL: INDENTATION HANDLING**
-When modifying indented code (like content inside tabs, functions, or loops), you MUST:
-- **Preserve exact indentation levels** in your added lines
-- **Show the complete indentation change** when moving code between indentation levels
-- **Include all whitespace** in your diff - indentation is part of the code structure
-- When unindenting code (removing tabs/context), show the original indented line with `-` and the unindented version with `+`
-- When indenting code (adding tabs/context), show the original unindented line with `-` and the indented version with `+`
-
-**HUNK HEADER FORMAT:**
-Use `@@ -START_LINE,1 +START_LINE,1 @@` where:
-- START_LINE is the line number in the **original file** where this hunk begins
-- Always use `1` for both count values (the system will calculate correct counts later)
-- All line numbers must reference the **original file**, not the modified version
-- For example, if the hunk begins on line 12, use `@@ -12,1 +12,1 @@`
-
-**WRONG FORMATS (DO NOT USE):**
-`@@ -12:` -> This is wrong because it is using a colon, doesn't have the count value, and doesn't have both sets of start_line numbers and lines counts.
-`@@ -12,1` -> This is wrong because it doesn't have both sets of start_line numbers and lines counts.
-`@@ 12,1 12,1 @@` -> This is wrong because it doesn't use - and + before the start_line 
-`@@-12,1 +12,1@@` -> This is wrong because it doesn't have a space after the first @@ and doesn't have a space before the second @@.
-
-**MULTIPLE HUNKS:**
-- If changes are separated by 5+ unchanged lines, create separate hunks
-- Each hunk needs its own `@@` header with the correct START_LINE for that section
-- Hunks must be in ascending order by line number
-
-<Example 1: Single change in middle of file>
-
-Assume `data_list = [` is on line 57 of the original file:
-```unified_diff
---- a/app.py 
-+++ b/app.py
-@@ -57,1 +57,1 @@
- data_list = [
--    {{'id': 1, 'name': 'Old'}},
-+    {{'id': 1, 'name': 'New'}},
-+    {{'id': 2, 'name': 'Also New'}},
+search_replace_instructions = f"""
+RESPONSE FORMAT: You can edit the existing code using the **SEARCH_REPLACE format** for exact string matching and replacement.
+
+**STRUCTURE:**
+```search_replace
+>>>>>>> SEARCH
+[exact code currently in the file]
+=======
+[new code to replace it with]
+<<<<<<< REPLACE
+```
+
+**COMPONENTS:**
+```search_replace - This is the start of the search/replace block
+- `>>>>>>> SEARCH` - Exact text that EXISTS NOW in the file (7 chevrons)
+- `=======` - Separator between the search and replace blocks (7 equals signs)  
+- `<<<<<<< REPLACE` - Replacement text (7 chevrons)
+
+---
+
+**CRITICAL RULES - READ CAREFULLY:**
+
+1. **SEARCH = CURRENT STATE ONLY**
+   - The SEARCH block must contain ONLY code that currently exists in the file
+   - NEVER include new code, future code, or code you wish existed in the SEARCH block
+   - Copy exact text from the current file, character-for-character
+   
+2. **EXACT MATCHING REQUIRED**
+   - Every space, tab, newline must match perfectly
+   - Preserve exact indentation (spaces vs tabs)
+   - Include trailing newlines if present
+   - No approximations - even one character difference will fail
+
+3. **SIZE LIMITS**
+   - There are no size limits to each search/replace block, however, it is generally preferable to keep the SEARCH blocks small and focused on one change. 
+   - For large changes, use multiple smaller search/replace blocks
+
+4. **UNIQUENESS**
+   - Include enough context to make the SEARCH block unique
+   - If text appears multiple times, add surrounding lines
+   - Ensure there's only ONE match in the file
+
+5. **VERIFICATION CHECKLIST** (before generating each block):
+   ✓ Is every line in my SEARCH block currently in the file?
+   ✓ Did I copy the exact spacing and whitespace?
+   ✓ Will this match exactly once?
+   
+6. **SEARCH REPLACE BLOCK STRUCTURE**
+   - You must adhere to to the exact search_replace structure as shown in the examples.
+
+---
+
+**MULTIPLE REPLACEMENTS:**
+- You can include multiple search/replace blocks in one response
+- Each block is independent and processed separately
+- Use separate ```search_replace blocks for each change
+
+<Example 1: Updating existing content>
+
+```search_replace
+>>>>>>> SEARCH
+st.title("Old Title")
+=======
+st.title("New Title")
+<<<<<<< REPLACE
 ```
 </Example 1>
 
-<Example 2: Multiple separate changes>
-Assume the original file has:
-
-Line 5: import os
-Line 30: def process():
-
-```unified_diff
---- a/app.py 
-+++ b/app.py
-@@ -5,1 +5,1 @@
- import os
-+import sys
-@@ -30,1 +30,1 @@
--def process():
-+def process_data():
+<Example 2: Adding new content>
+
+```search_replace
+>>>>>>> SEARCH
+st.title("My App")
+=======
+st.title("My App")
+st.header("Welcome")
+st.write("This is a test app")
+<<<<<<< REPLACE
 ```
 </Example 2>
 
-<Example 3: Adding multiple entries to a list while respecting indentations>
-
-In the example below, assume that the line of code `data_list = [` is on line 57 of the existing streamlit app.
-
-```unified_diff
---- a/app.py 
-+++ b/app.py
-@@ -57,1 +57,1 @@
- data_list = [
-     {{'id': 1, 'name': 'Item A', 'category': 'Type 1', 'value': 100}},
-     {{'id': 2, 'name': 'Item B', 'category': 'Type 2', 'value': 200}},
--    {MITO_TODO_PLACEHOLDER}: Add remaining entries from notebook
-+    {{'id': 3, 'name': 'Item C', 'category': 'Type 3', 'value': 300}},
-+    {{'id': 4, 'name': 'Item D', 'category': 'Type 4', 'value': 400}},
-+    {{'id': 5, 'name': 'Item E', 'category': 'Type 5', 'value': 500}},
-+    {{'id': 6, 'name': 'Item F', 'category': 'Type 6', 'value': 600}},
-+    {{'id': 7, 'name': 'Item G', 'category': 'Type 7', 'value': 700}},
-+    {{'id': 8, 'name': 'Item H', 'category': 'Type 8', 'value': 800}},
-+    {{'id': 9, 'name': 'Item I', 'category': 'Type 9', 'value': 900}},
-+    {{'id': 10, 'name': 'Item J', 'category': 'Type 10', 'value': 1000}}
+<Example 3: Deleting existing content>
+
+```search_replace
+>>>>>>> SEARCH
+st.write("Old message")
+=======
+<<<<<<< REPLACE
 ```
 </Example 3>
 
-<Example 4: Consolidating tabs - removing tab structure and unindenting content>
+<Example 4: Multiple replacements in one response>
+
+```search_replace
+>>>>>>> SEARCH
+st.title("Old Title")
+=======
+st.title("New Title")
+<<<<<<< REPLACE
+```
+
+```search_replace
+>>>>>>> SEARCH
+st.write("Old message")
+=======
+st.write("New message")
+<<<<<<< REPLACE
+```
+</Example 4>
+
+<Example 5: Using extra context to identify the correct code to replace>
+
+In the below example, assume that the code st.write("Old message") appears multiple times in the file, so we use extra context lines to identify the correct code to replace.
+
+```search_replace
+>>>>>>> SEARCH
+# This is a unique comment
+st.write("Old message")
+=======
+# This is a unique comment
+st.write("New message")
+<<<<<<< REPLACE
+```
+</Example 5>
+
+<Example 6: Search/replace while respecting whitespace and indentation>
+
+```search_replace
+>>>>>>> SEARCH
+data_list = [
+    {{'id': 1, 'name': 'Item A'}},
+    {MITO_TODO_PLACEHOLDER}: Add remaining entries from notebook
+]
+=======
+data_list = [
+    {{'id': 1, 'name': 'Item A'}},
+    {{'id': 2, 'name': 'Item B'}},
+    {{'id': 3, 'name': 'Item C'}},
+    {{'id': 4, 'name': 'Item D'}}
+]
+<<<<<<< REPLACE
+```
+</Example 6>
+
+<Example 7: Tab structure changes>
 
-Assume the original file has tabs starting at line 10:
-```python
+```search_replace
+>>>>>>> SEARCH
 tab1, tab2 = st.tabs(["Cat", "Dog"])
 
 with tab1:
@@ -126,57 +159,14 @@ with tab1:
 with tab2:
     st.header("A dog")
     st.image("https://static.streamlit.io/examples/dog.jpg", width=200)
-```
-
-To consolidate into a single screen without tabs:
-```unified_diff
---- a/app.py 
-+++ b/app.py
-@@ -10,1 +10,1 @@
--tab1, tab2 = st.tabs(["Cat", "Dog"])
--
--with tab1:
--    st.header("A cat")
--    st.image("https://static.streamlit.io/examples/cat.jpg", width=200)
--with tab2:
--    st.header("A dog")
--    st.image("https://static.streamlit.io/examples/dog.jpg", width=200)
-+st.header("A cat")
-+st.image("https://static.streamlit.io/examples/cat.jpg", width=200)
-+st.header("A dog")
-+st.image("https://static.streamlit.io/examples/dog.jpg", width=200)
-```
-</Example 4>
-
-<Example 5: Adding tab structure - indenting existing content>
-
-Assume the original file has content starting at line 10:
-```python
+=======
 st.header("A cat")
 st.image("https://static.streamlit.io/examples/cat.jpg", width=200)
 st.header("A dog")
 st.image("https://static.streamlit.io/examples/dog.jpg", width=200)
+<<<<<<< REPLACE
 ```
+</Example 7>
 
-To add tab structure:
-```unified_diff
---- a/app.py 
-+++ b/app.py
-@@ -10,1 +10,1 @@
--st.header("A cat")
--st.image("https://static.streamlit.io/examples/cat.jpg", width=200)
--st.header("A dog")
--st.image("https://static.streamlit.io/examples/dog.jpg", width=200)
-+tab1, tab2 = st.tabs(["Cat", "Dog"])
-+
-+with tab1:
-+    st.header("A cat")
-+    st.image("https://static.streamlit.io/examples/cat.jpg", width=200)
-+with tab2:
-+    st.header("A dog")
-+    st.image("https://static.streamlit.io/examples/dog.jpg", width=200)
-```
-</Example 5>
-
-Your response must consist **only** of valid unified-diff block.
+Your response must consist **only** of valid search_replace blocks.
 """