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

mito_ai/streamlit_conversion/agent_utils.py

@@ -3,8 +3,14 @@
 
 from typing import List
 import re
+from anthropic.types import MessageParam
+from mito_ai.streamlit_conversion.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"
 
 def extract_todo_placeholders(agent_response: str) -> List[str]:
     """Extract TODO placeholders from the agent's response"""
@@ -23,6 +29,8 @@ def apply_patch_to_text(text: str, diff: str) -> str:
     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
     -------
@@ -42,37 +50,46 @@ def apply_patch_to_text(text: str, diff: str) -> str:
     patch = PatchSet(diff.splitlines(keepends=True))
 
     # We expect a single-file patch (what the prompt asks the model to emit)
-    if len(patch) != 1:
+    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 a patch for exactly one file, got {len(patch)} files."
+            f"Expected patches for exactly one file, got files: {file_names}"
         )
 
-    file_patch = patch[0]
-
+    # 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)
 
-    for hunk in file_patch:
-        # Copy unchanged lines before this hunk
-        while cursor < hunk.source_start - 1:
-            result_lines.append(original_lines[cursor])
-            cursor += 1
-
-        # Apply hunk line-by-line
-        for line in hunk:
-            if line.is_context:
-                result_lines.append(original_lines[cursor])
+    # 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
-            elif line.is_removed:
-                cursor += 1  # Skip this line from the original
-            elif line.is_added:
-                # Ensure added line ends with newline for consistency
-                val = line.value
-                if not val.endswith("\n"):
-                    val += "\n"
-                result_lines.append(val)
+
+            # 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:])
@@ -80,11 +97,66 @@ def apply_patch_to_text(text: str, diff: str) -> str:
     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):
@@ -99,18 +171,64 @@ def fix_diff_headers(diff: str) -> str:
                 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)):
-                    next_line = lines[j]
-                    if next_line.startswith('@@') or next_line.startswith('---'):
+                    if lines[j].startswith('@@'):
+                        hunk_end = j
                         break
-                    if next_line.startswith(' ') or next_line.startswith('-'):
+                
+                # 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 next_line.startswith(' ') or next_line.startswith('+'):
+                    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} @@"
     
-    return '\n'.join(lines)
-
-
+    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
+    temperature = 0.2
+
+    accumulated_response = ""
+    async for stream_chunk in stream_anthropic_completion_from_mito_server(
+        model = model,
+        max_tokens = max_tokens,
+        temperature = temperature,
+        system = streamlit_system_prompt,
+        messages = message_to_agent,
+        stream=True,
+        message_type=MessageType.STREAMLIT_CONVERSION,
+        reply_fn=None,
+        message_id=""
+    ):
+        accumulated_response += stream_chunk
+    return accumulated_response

mito_ai/streamlit_conversion/prompts/prompt_constants.py

@@ -3,26 +3,100 @@
 
 MITO_TODO_PLACEHOLDER = "# MITO_TODO_PLACEHOLDER"
 
-unified_diff_instrucrions = f"""
+unified_diff_instructions = f"""
 RESPONSE FORMAT: Return the changes you want to make to the streamlit app as a **unified diff (git-style patch)**:
-- Begin with a ````unified_diff` header and a ```` end header.
-- Then, include the standard header lines `--- a/app.py` and `+++ b/app.py`.
+
+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.
+  * 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.
 
-**IMPORTANT: For the hunk header, use `@@ -START_LINE,1 +START_LINE,1 @@` where we always use 1 as the count value. In a later step, the system will automatically calculate the correct counts.**
+**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'}},
+```
+</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>
 
-<Example Response>
+<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
+--- a/app.py 
 +++ b/app.py
 @@ -57,1 +57,1 @@
  data_list = [
@@ -36,24 +110,73 @@ In the example below, assume that the line of code `data_list = [` is on line 57
 +    {{'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}},
-+    {{'id': 11, 'name': 'Item K', 'category': 'Type 11', 'value': 1100}},
-+    {{'id': 12, 'name': 'Item L', 'category': 'Type 12', 'value': 1200}},
-+    {{'id': 13, 'name': 'Item M', 'category': 'Type 13', 'value': 1300}},
-+    {{'id': 14, 'name': 'Item N', 'category': 'Type 14', 'value': 1400}},
-+    {{'id': 15, 'name': 'Item O', 'category': 'Type 15', 'value': 1500}},
-+    {{'id': 16, 'name': 'Item P', 'category': 'Type 16', 'value': 1600}},
-+    {{'id': 17, 'name': 'Item Q', 'category': 'Type 17', 'value': 1700}},
-+    {{'id': 18, 'name': 'Item R', 'category': 'Type 18', 'value': 1800}},
-+    {{'id': 19, 'name': 'Item S', 'category': 'Type 19', 'value': 1900}},
-+    {{'id': 20, 'name': 'Item T', 'category': 'Type 20', 'value': 2000}},
-+    {{'id': 21, 'name': 'Item U', 'category': 'Type 21', 'value': 2100}},
-+    {{'id': 22, 'name': 'Item V', 'category': 'Type 22', 'value': 2200}},
-+    {{'id': 23, 'name': 'Item W', 'category': 'Type 23', 'value': 2300}},
-+    {{'id': 24, 'name': 'Item X', 'category': 'Type 24', 'value': 2400}},
-+    {{'id': 25, 'name': 'Item Y', 'category': 'Type 25', 'value': 2500}}
++    {{'id': 10, 'name': 'Item J', 'category': 'Type 10', 'value': 1000}}
+```
+</Example 3>
+
+<Example 4: Consolidating tabs - removing tab structure and unindenting content>
+
+Assume the original file has tabs starting at line 10:
+```python
+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)
+```
+
+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)
+```
+
+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 Response>
+</Example 5>
 
 Your response must consist **only** of valid unified-diff block.
 """

mito_ai/streamlit_conversion/prompts/streamlit_app_creation_prompt.py

@@ -1,9 +1,10 @@
 # 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 MITO_TODO_PLACEHOLDER
 
-def get_streamlit_app_creation_prompt(notebook: dict) -> str:
+def get_streamlit_app_creation_prompt(notebook: List[dict]) -> str:
     """
     This prompt is used to create a streamlit app from a notebook.
     """

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_instrucrions
+from mito_ai.streamlit_conversion.prompts.prompt_constants import unified_diff_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_instrucrions}
+{unified_diff_instructions}
 
 ===============================================
 

mito_ai/streamlit_conversion/prompts/streamlit_finish_todo_prompt.py

@@ -1,10 +1,11 @@
 # 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 MITO_TODO_PLACEHOLDER, unified_diff_instrucrions
+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_utils import add_line_numbers_to_code
 
-def get_finish_todo_prompt(notebook: dict, existing_streamlit_app_code: str, todo_placeholder: str) -> str:
+def get_finish_todo_prompt(notebook: List[dict], existing_streamlit_app_code: str, todo_placeholder: str) -> str:
   
     existing_streamlit_app_code_with_line_numbers = add_line_numbers_to_code(existing_streamlit_app_code)
     
@@ -24,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_instrucrions}
+{unified_diff_instructions}
 
 ===============================================
 

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 unified_diff_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.
+
+{unified_diff_instructions}
+
+===============================================
+
+INPUT NOTEBOOK: 
+{notebook}
+
+===============================================
+
+EXISTING STREAMLIT APP:
+{existing_streamlit_app_code_with_line_numbers}
+
+===============================================
+
+USER EDIT REQUEST:
+{edit_prompt}
+
+"""