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

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.
 """

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}
 
 ===============================================