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

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}
+
+"""

mito_ai/streamlit_conversion/streamlit_agent_handler.py

@@ -1,137 +1,138 @@
 # Copyright (c) Saga Inc.
 # Distributed under the terms of the GNU Affero General Public License v3.0 License.
 
-import logging
 import os
 from anthropic.types import MessageParam
 from typing import List, Optional, Tuple, cast
-
-from mito_ai.logger import get_logger
-from mito_ai.streamlit_conversion.agent_utils import apply_patch_to_text, extract_todo_placeholders, fix_diff_headers
+from mito_ai.streamlit_conversion.agent_utils import apply_patch_to_text, extract_todo_placeholders, fix_diff_headers, 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.streamlit_system_prompt import streamlit_system_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, parse_jupyter_notebook_to_extract_required_content
-from mito_ai.utils.anthropic_utils import stream_anthropic_completion_from_mito_server
+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.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
 
-STREAMLIT_AI_MODEL = "claude-3-5-haiku-latest"
-
-class StreamlitCodeGeneration:
-    @property
-    def log(self) -> logging.Logger:
-        """Use Mito AI logger."""
-        return get_logger()
-
-    async def get_response_from_agent(self, 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
-
-        self.log.info("Getting response from agent...")
-        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
-
-    async def generate_streamlit_code(self, notebook: dict) -> str:
-        """Send a query to the agent, get its response and parse the code"""
-        
-        prompt_text = get_streamlit_app_creation_prompt(notebook)
-        
-        messages: List[MessageParam] = [
-            cast(MessageParam, {
-                "role": "user",
-                "content": [{
-                    "type": "text",
-                    "text": prompt_text
-                }]
-            })
-        ]
-        
-        agent_response = await self.get_response_from_agent(messages)
-        converted_code = extract_code_blocks(agent_response)
-        
-        # Extract the TODOs from the agent's response
-        todo_placeholders = extract_todo_placeholders(agent_response)
-        
-        for todo_placeholder in todo_placeholders:
-            print(f"Processing AI TODO: {todo_placeholder}")
-            todo_prompt = get_finish_todo_prompt(notebook, converted_code, todo_placeholder)
-            todo_messages: List[MessageParam] = [
-                cast(MessageParam, {
-                    "role": "user",
-                    "content": [{
-                        "type": "text",
-                        "text": todo_prompt
-                    }]
-                })
-            ]
-            todo_response = await self.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)
-                    
-        return converted_code
-
+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
 
-    async def correct_error_in_generation(self, error: str, streamlit_app_code: str) -> str:
-        """If errors are present, send it back to the agent to get corrections in code"""
-        messages: List[MessageParam] = [
+async def generate_new_streamlit_code(notebook: List[dict]) -> str:
+    """Send a query to the agent, get its response and parse the code"""
+    
+    prompt_text = get_streamlit_app_creation_prompt(notebook)
+    
+    messages: List[MessageParam] = [
+        cast(MessageParam, {
+            "role": "user",
+            "content": [{
+                "type": "text",
+                "text": prompt_text
+            }]
+        })
+    ]
+    agent_response = await get_response_from_agent(messages)
+    converted_code = extract_code_blocks(agent_response)
+    
+    # Extract the TODOs from the agent's response
+    todo_placeholders = extract_todo_placeholders(agent_response)
+    
+    for todo_placeholder in todo_placeholders:
+        print(f"Processing AI TODO: {todo_placeholder}")
+        todo_prompt = get_finish_todo_prompt(notebook, converted_code, todo_placeholder)
+        todo_messages: List[MessageParam] = [
             cast(MessageParam, {
                 "role": "user",
                 "content": [{
                     "type": "text",
-                    "text": get_streamlit_error_correction_prompt(error, streamlit_app_code)
+                    "text": todo_prompt
                 }]
             })
         ]
-        agent_response = await self.get_response_from_agent(messages)
+        todo_response = await get_response_from_agent(todo_messages)
         
         # Apply the diff to the streamlit app
-        exctracted_diff = extract_unified_diff_blocks(agent_response)
-        
-        print(f"\n\nExtracted diff: {exctracted_diff}")
+        exctracted_diff = extract_unified_diff_blocks(todo_response)
         fixed_diff = fix_diff_headers(exctracted_diff)
-        streamlit_app_code = apply_patch_to_text(streamlit_app_code, fixed_diff)
-        
-        print("\n\nUpdated app code: ", streamlit_app_code)
+        converted_code = apply_patch_to_text(converted_code, fixed_diff)
+                
+    return converted_code
 
-        return streamlit_app_code
 
+async def update_existing_streamlit_code(notebook: List[dict], streamlit_app_code: str, edit_prompt: str) -> str:
+    """Send a query to the agent, get its response and parse the code"""
+    prompt_text = get_update_existing_app_prompt(notebook, streamlit_app_code, edit_prompt)
+    
+    messages: List[MessageParam] = [
+        cast(MessageParam, {
+            "role": "user",
+            "content": [{
+                "type": "text",
+                "text": prompt_text
+            }]
+        })
+    ]
+    
+    agent_response = await get_response_from_agent(messages)
+    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)
+    return converted_code
+
+
+async def correct_error_in_generation(error: str, streamlit_app_code: str) -> str:
+    """If errors are present, send it back to the agent to get corrections in code"""
+    messages: List[MessageParam] = [
+        cast(MessageParam, {
+            "role": "user",
+            "content": [{
+                "type": "text",
+                "text": get_streamlit_error_correction_prompt(error, streamlit_app_code)
+            }]
+        })
+    ]
+    agent_response = await get_response_from_agent(messages)
+    
+    # Apply the 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)
+
+    return streamlit_app_code
 
-async def streamlit_handler(notebook_path: str) -> Tuple[bool, Optional[str], str]:
+async def streamlit_handler(notebook_path: str, edit_prompt: str = "") -> Tuple[bool, Optional[str], str]:
     """Handler function for streamlit code generation and validation"""
 
     clean_directory_check(notebook_path)
 
     notebook_code = parse_jupyter_notebook_to_extract_required_content(notebook_path)
-    streamlit_code_generator = StreamlitCodeGeneration()
-
-    streamlit_code = await streamlit_code_generator.generate_streamlit_code(notebook_code)
+    app_directory = get_app_directory(notebook_path)
     
+    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)
+        
+        if streamlit_code is None:
+            return False, '', "Error updating existing streamlit app because app.py file was not found."
+        
+        streamlit_code = await update_existing_streamlit_code(notebook_code, streamlit_code, edit_prompt)
+    else:
+        # Otherwise generate a new streamlit app
+        streamlit_code = await generate_new_streamlit_code(notebook_code)
+       
+    # Then, after creating/updating the app, validate that the new code runs 
     has_validation_error, errors = validate_app(streamlit_code, notebook_path)
     tries = 0
     while has_validation_error and tries < 5:
         for error in errors:
-            streamlit_code = await streamlit_code_generator.correct_error_in_generation(error, streamlit_code)
+            streamlit_code = await correct_error_in_generation(error, streamlit_code)
         
         has_validation_error, errors = validate_app(streamlit_code, notebook_path)
         
@@ -142,20 +143,13 @@ async def streamlit_handler(notebook_path: str) -> Tuple[bool, Optional[str], st
         tries+=1
 
     if has_validation_error:
-        log_streamlit_app_creation_error('mito_server_key', MessageType.STREAMLIT_CONVERSION, error)
+        log_streamlit_app_creation_error('mito_server_key', MessageType.STREAMLIT_CONVERSION, error, edit_prompt)
         return False, '', "Error generating streamlit code by agent"
     
-    # Convert to absolute path for directory calculation
-    absolute_notebook_path = notebook_path
-    if not (notebook_path.startswith('/') or (len(notebook_path) > 1 and notebook_path[1] == ':')):
-        absolute_notebook_path = os.path.join(os.getcwd(), notebook_path)
-    
-    app_directory = os.path.dirname(absolute_notebook_path)
-    
+    # 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)
+        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)
+    log_streamlit_app_creation_success('mito_server_key', MessageType.STREAMLIT_CONVERSION, edit_prompt)
     return success_flag, app_path, message

mito_ai/streamlit_conversion/streamlit_system_prompt.py

@@ -22,6 +22,7 @@ STREAMLIT IMPLEMENTATION GUIDELINES:
 - Include all text explanations and insights from markdown cells
 - Add interactive elements where beneficial (filters, selectors, etc.)
 - Ensure professional styling and layout suitable for executives
+- Just create the streamlit app code, do not include a _main_ function block. The file will be run directly using `streamlit run app.py`.
 
 CRITICAL REQUIREMENTS:
 1. **PRESERVE ALL CODE EXACTLY**: Every line of code, every data structure, every import must be included in full

mito_ai/streamlit_conversion/streamlit_utils.py

@@ -4,7 +4,7 @@
 import re
 import json
 import os
-from typing import Dict, Optional, Tuple, Any
+from typing import Dict, List, Optional, Tuple, Any
 from pathlib import Path
 
 def extract_code_blocks(message_content: str) -> str:
@@ -64,6 +64,13 @@ def create_app_file(app_directory: str, code: str) -> Tuple[bool, str, str]:
     except Exception as e:
         return False, '', f"Unexpected error: {str(e)}"
     
+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
+    with open(app_path, 'r', encoding='utf-8') as f:
+        return f.read()
+    
 
 def get_app_path(app_directory: str) -> Optional[str]:
     """
@@ -73,9 +80,8 @@ def get_app_path(app_directory: str) -> Optional[str]:
     if not os.path.exists(app_path):
         return None
     return app_path
-    
 
-def parse_jupyter_notebook_to_extract_required_content(notebook_path: str) -> Dict[str, Any]:
+def parse_jupyter_notebook_to_extract_required_content(notebook_path: str) -> List[Dict[str, Any]]:
     """
     Read a Jupyter notebook and filter cells to keep only cell_type and source fields.
 
@@ -105,18 +111,15 @@ def parse_jupyter_notebook_to_extract_required_content(notebook_path: str) -> Di
             raise KeyError("Notebook does not contain 'cells' key")
 
         # Filter each cell to keep only cell_type and source
-        filtered_cells = []
+        filtered_cells: List[Dict[str, Any]] = []
         for cell in notebook_data['cells']:
-            filtered_cell = {
+            filtered_cell: Dict[str, Any] = {
                 'cell_type': cell.get('cell_type', ''),
                 'source': cell.get('source', [])
             }
             filtered_cells.append(filtered_cell)
-
-        # Update the notebook data with filtered cells
-        notebook_data['cells'] = filtered_cells
-
-        return notebook_data
+            
+        return filtered_cells
 
     except FileNotFoundError:
         raise FileNotFoundError(f"Notebook file not found: {notebook_path}")