mito-ai 0.1.53__py3-none-any.whl → 0.1.55__py3-none-any.whl

mito_ai/_version.py

@@ -1,4 +1,4 @@
 # This file is auto-generated by Hatchling. As such, do not:
 #   - modify
 #   - track in version control e.g. be sure to add to .gitignore
-__version__ = VERSION = '0.1.53'
+__version__ = VERSION = '0.1.55'

mito_ai/completions/models.py

@@ -35,7 +35,7 @@ class AgentResponse(BaseModel):
     get_cell_output_cell_id: Optional[str]
     next_steps: Optional[List[str]]
     analysis_assumptions: Optional[List[str]]
-    edit_streamlit_app_prompt: Optional[str]
+    streamlit_app_prompt: Optional[str]
     
     
 @dataclass(frozen=True)

mito_ai/completions/prompt_builders/agent_system_message.py

@@ -3,6 +3,7 @@
 
 from mito_ai.completions.prompt_builders.prompt_constants import (
     CITATION_RULES,
+    CELL_REFERENCE_RULES,
     FILES_SECTION_HEADING,
     JUPYTER_NOTEBOOK_SECTION_HEADING,
     VARIABLES_SECTION_HEADING,
@@ -231,15 +232,17 @@ When you want to create a new Streamlit app from the current notebook, respond w
 
 {{
     type: 'create_streamlit_app',
+    streamlit_app_prompt: str
     message: str
 }}
 
 Important information:
-1. The message is a short summary of why you're creating the Streamlit app.
-2. Only use this tool when the user explicitly asks to create or preview a Streamlit app AND no Streamlit app is currently open.
-3. This tool creates a new app from scratch - use EDIT_STREAMLIT_APP tool if the user is asking you to edit, update, or modify an app that already exists.
-4. Using this tool will automatically open the app so the user can see a preview of the app.
-5. When you use this tool, assume that it successfully created the Streamlit app unless the user explicitly tells you otherwise. The app will remain open so that the user can view it until the user decides to close it. You do not need to continually use the create_streamlit_app tool to keep the app open.
+1. The streamlit_app_prompt is a short description of how the app should be structured. It should be a high level specification that includes things like what fields should be configurable, what tabs should exist, etc. It does not need to be overly detailed however.
+2. The message is a short summary of why you're creating the Streamlit app.
+3. Only use this tool when the user explicitly asks to create or preview a Streamlit app. If the streamlit app for this app already exists, then use an empty string '' as the streamlit_app_prompt.
+4. This tool creates a new app from scratch - use EDIT_STREAMLIT_APP tool if the user is asking you to edit, update, or modify an app that already exists.
+5. Using this tool will automatically open the app so the user can see a preview of the app. If the user is asking you to open an app that already exists, but not make any changes to the app, this is the correct tool.
+6. When you use this tool, assume that it successfully created the Streamlit app unless the user explicitly tells you otherwise. The app will remain open so that the user can view it until the user decides to close it. You do not need to continually use the create_streamlit_app tool to keep the app open.
 
 <Example>
 
@@ -248,6 +251,7 @@ Your task: Show me my notebook as an app.
 Output:
 {{
     type: 'create_streamlit_app',
+    streamlit_app_prompt: "The app should have a beginning date and end date input field at the top. It should then be followed by two tabs for the user to select between: current performance and projected performance.",
     message: "I'll convert your notebook into an app."
 }}
 
@@ -264,12 +268,12 @@ When you want to edit an existing Streamlit app, respond with this format:
 {{
     type: 'edit_streamlit_app',
     message: str,
-    edit_streamlit_app_prompt: str
+    streamlit_app_prompt: str
 }}
 
 Important information:
 1. The message is a short summary of why you're editing the Streamlit app.
-2. The edit_streamlit_app_prompt is REQUIRED and must contain specific instructions for the edit (e.g., "Make the title text larger", "Change the chart colors to blue", "Add a sidebar with filters").
+2. The streamlit_app_prompt is REQUIRED and must contain specific instructions for the edit (e.g., "Make the title text larger", "Change the chart colors to blue", "Add a sidebar with filters").
 3. Only use this tool when the user asks to edit, update, or modify a Streamlit app. 
 4. The app does not need to already be open for you to use the tool. Using this tool will automatically open the streamlit app after applying the changes so the user can view it. You do not need to call the create_streamlit_app tool first.
 5. When you use this tool, assume that it successfully edited the Streamlit app unless the user explicitly tells you otherwise. The app will remain open so that the user can view it until the user decides to close it. 
@@ -288,7 +292,7 @@ When you have completed the user's task, respond with a message in this format:
 
 Important information:
 1. The message is a short summary of the ALL the work that you've completed on this task. It should not just refer to the final message. It could be something like "I've completed the sales strategy analysis by exploring key relationships in the data and summarizing creating a report with three recommendations to boost sales.""
-2. The message should include citations for any insights that you shared with the user.
+2. The message should include citations for any insights that you shared with the user and cell references for whenever you refer to specific cells that you've updated or created.
 3. The next_steps is an optional list of 2 or 3 suggested follow-up tasks or analyses that the user might want to perform next. These should be concise, actionable suggestions that build on the work you've just completed. For example: ["Export the cleaned data to CSV", "Analyze revenue per customer", "Convert notebook into an app"].
 4. The next_steps should be as relevant to the user's actual task as possible. Try your best not to make generic suggestions like "Analyze the data" or "Visualize the results". For example, if the user just asked you to calculate LTV of their customers, you might suggest the following next steps: ["Graph key LTV drivers: churn and average transaction value", "Visualize LTV per age group"].
 5. If you are not sure what the user might want to do next, err on the side of suggesting next steps instead of making an assumption and using more CELL_UPDATES.
@@ -322,19 +326,29 @@ Output:
 
 RULES
 
-- You are working in a Jupyter Lab environment in a .ipynb file. 
+- You are working in a Jupyter Lab environment in a .ipynb file.
 - You can only respond with CELL_UPDATES or FINISHED_TASK responses.
 - In each message you send to the user, you can send one CellModification, one CellAddition, or one FINISHED_TASK response. BUT YOU WILL GET TO SEND MULTIPLE MESSAGES TO THE USER TO ACCOMPLISH YOUR TASK SO DO NOT TRY TO ACCOMPLISH YOUR TASK IN A SINGLE MESSAGE.
 - After you send a CELL_UPDATE, the user will send you a message with the updated variables, code, and files in the current directory. You will use this information to decide what to do next, so it is critical that you wait for the user's response after each CELL_UPDATE before deciding your next action.
 - When updating code, keep as much of the original code as possible and do not recreate variables that already exist.
-- When you want to display a dataframe to the user, just write the dataframe on the last line of the code cell instead of writing print(<dataframe name>). Jupyter will automatically display the dataframe in the notebook.
 - When writing the message, do not explain to the user how to use the CELL_UPDATE or FINISHED_TASK response, they will already know how to use them. Just provide a summary of your thought process. Do not reference any Cell IDs in the message.
 - When writing the message, do not include leading words like "Explanation:" or "Thought process:". Just provide a summary of your thought process.
 - When writing the message, use tickmarks when referencing specific variable names. For example, write `sales_df` instead of "sales_df" or just sales_df.
 
+====
+
+CODE STYLE
+
+- Avoid using try/except blocks and other defensive programming patterns (like checking if files exist before reading them, verifying variables are defined before using them, etc.) unless there is a really good reason. In Jupyter notebooks, errors should surface immediately so users can identify and fix issues. When errors are caught and suppressed or when defensive checks hide problems, users continue running broken code without realizing it, and the agent's auto-error-fix loop cannot trigger. If a column doesn't exist, a file is missing, a variable isn't defined, or a module isn't installed, let it error. The user needs to know.
+- When you want to display a dataframe to the user, just write the dataframe on the last line of the code cell instead of writing print(<dataframe name>). Jupyter will automatically display the dataframe in the notebook.
+- When importing matplotlib, write the code `%matplotlib inline` to make sure the graphs render in Jupyter.
+
 ==== 
 {CITATION_RULES}
 
+====
+{CELL_REFERENCE_RULES}
+
 <Citation Example>
 
 ### User Message 1:
@@ -468,5 +482,4 @@ REMEMBER, YOU ARE GOING TO COMPLETE THE USER'S TASK OVER THE COURSE OF THE ENTIR
 ====
 
 OTHER USEFUL INFORMATION:
-1. When importing matplotlib, write the code `%matplotlib inline` to make sure the graphs render in Jupyter
-2. The active cell ID is shared with you so that when the user refers to "this cell" or similar phrases, you know which cell they mean. However, you are free to edit any cell that you see fit."""
+1. The active cell ID is shared with you so that when the user refers to "this cell" or similar phrases, you know which cell they mean. However, you are free to edit any cell that you see fit."""

mito_ai/completions/prompt_builders/chat_system_message.py

@@ -3,7 +3,8 @@
 
 from mito_ai.completions.prompt_builders.prompt_constants import (
     CHAT_CODE_FORMATTING_RULES,
-    CITATION_RULES, 
+    CITATION_RULES,
+    CELL_REFERENCE_RULES,
     ACTIVE_CELL_ID_SECTION_HEADING, 
     CODE_SECTION_HEADING,
     get_database_rules
@@ -28,6 +29,9 @@ Other useful information:
 ====
 {CITATION_RULES}
 
+====
+{CELL_REFERENCE_RULES}
+
 <Example 1> 
 {ACTIVE_CELL_ID_SECTION_HEADING}
 '7b3a9e2c-5d14-4c83-b2f9-d67891e4a5f2'
@@ -79,6 +83,18 @@ Notice in the example above that the user is just sending a friendly message, so
 ==== 
 {CHAT_CODE_FORMATTING_RULES}
 
+====
+
+CODE STYLE
+
+- Avoid using try/except blocks and other defensive programming patterns (like checking if files exist before reading them, verifying variables are defined before using them, etc.) unless there is a really good reason. In Jupyter notebooks, errors should surface immediately so users can identify and fix issues. When errors are caught and suppressed or when defensive checks hide problems, users continue running broken code without realizing it, and the agent's auto-error-fix loop cannot trigger. If a column doesn't exist, a file is missing, a variable isn't defined, or a module isn't installed, let it error. The user needs to know.
+- Write code that preserves the intent of the original code shared with you and the task to complete.
+- Make the solution as simple as possible.
+- Do not add temporary comments like '# Fixed the typo here' or '# Added this line to fix the error'
+- When importing matplotlib, write the code `%matplotlib inline` to make sure the graphs render in Jupyter.
+
+====
+
 IMPORTANT RULES:
 - Do not recreate variables that already exist
 - Keep as much of the original code as possible
@@ -87,6 +103,5 @@ IMPORTANT RULES:
 - Write code that preserves the intent of the original code shared with you and the task to complete.
 - Make the solution as simple as possible.
 - Reuse as much of the existing code as possible.
-- Do not add temporary comments like '# Fixed the typo here' or '# Added this line to fix the error'
 - Whenever writing Python code, it should be a python code block starting with ```python and ending with ```
 """

mito_ai/completions/prompt_builders/prompt_constants.py

@@ -46,6 +46,28 @@ Citation Rules:
 8. Do not include the citation in the code block as a comment. ONLY include the citation in the message field of your response.
 """
 
+CELL_REFERENCE_RULES = """RULES FOR REFERENCING CELLS
+
+When referring to specific cells in the notebook in your messages, use cell references so the user can easily navigate to the cell you're talking about. The user sees cells numbered as "Cell 1", "Cell 2", etc., but internally cells are identified by their unique IDs.
+
+To reference a cell, use this format inline in your message:
+[MITO_CELL_REF:cell_id]
+
+This will be displayed to the user as a clickable "Cell N" link that navigates to the referenced cell.
+
+Cell Reference Rules:
+
+1. Use cell references when discussing specific cells you've created or modified (e.g., "I've added the data cleaning code in [MITO_CELL_REF:abc123]").
+2. Use cell references when referring to cells the user mentioned or that contain relevant context.
+3. The cell_id must be an actual cell ID from the notebook - do not make up IDs.
+4. Place the reference inline where it makes sense in your message, similar to how you would write "Cell 3" in natural language.
+5. Do not use cell references in code - only in the message field of your responses.
+6. Cell references are different from citations. Use citations for specific line-level insights; use cell references for general cell-level navigation.
+
+Example:
+"I've loaded the sales data in [MITO_CELL_REF:c68fdf19-db8c-46dd-926f-d90ad35bb3bc] and will now calculate the monthly totals."
+"""
+
 def get_active_cell_output_str(has_active_cell_output: bool) -> str:
     """
     Used to tell the AI about the output of the active code cell. 

mito_ai/completions/prompt_builders/utils.py

@@ -39,6 +39,7 @@ def get_selected_context_str(additional_context: Optional[List[Dict[str, str]]])
     selected_files = [context["value"] for context in additional_context if context.get("type") == "file"]
     selected_db_connections = [context["value"] for context in additional_context if context.get("type") == "db"]
     selected_images = [context["value"] for context in additional_context if context.get("type", "").startswith("image/")]
+    selected_cells = [context["value"] for context in additional_context if context.get("type") == "cell"]
 
     # STEP 2: Create a list of strings (instructions) for each context type
     context_parts = []
@@ -66,6 +67,12 @@ def get_selected_context_str(additional_context: Optional[List[Dict[str, str]]])
             "The following images have been selected by the user to be used in the task:\n"
             + "\n".join(selected_images)
         )
+        
+    if len(selected_cells) > 0:
+        context_parts.append(
+            "The following cells have been selected by the user to be used in the task:\n"
+            + "\n".join(selected_cells)
+        )
 
     # STEP 3: Combine into a single string
     return "\n\n".join(context_parts)

mito_ai/streamlit_conversion/prompts/streamlit_app_creation_prompt.py

@@ -4,10 +4,24 @@
 from typing import List
 from mito_ai.streamlit_conversion.prompts.prompt_constants import MITO_TODO_PLACEHOLDER
 
-def get_streamlit_app_creation_prompt(notebook: List[dict]) -> str:
+def get_streamlit_app_spec_section(streamlit_app_prompt: str) -> str:
+    if streamlit_app_prompt == '':
+        return ''
+    
+    return f"""
+Here is a high level outline of the streamlit app. Use your best judgement to implement this structure.
+
+{streamlit_app_prompt}
+
+"""
+    
+
+def get_streamlit_app_creation_prompt(notebook: List[dict], streamlit_app_prompt: str) -> str:
     """
     This prompt is used to create a streamlit app from a notebook.
     """
+    streamlit_app_spec_section = get_streamlit_app_spec_section(streamlit_app_prompt)
+    
     return f"""Convert the following Jupyter notebook into a Streamlit application.
 
 GOAL: Create a complete, runnable Streamlit app that accurately represents the notebook. It must completely convert the notebook. 
@@ -40,7 +54,9 @@ data = [
 ]
 </Example>
 
-Notebook to convert:
+{streamlit_app_spec_section}
+
+NOTEBOOK TO CONVERT:
 
 {notebook}
 """

mito_ai/streamlit_conversion/streamlit_agent_handler.py

@@ -16,10 +16,10 @@ from mito_ai.utils.error_classes import StreamlitConversionError
 from mito_ai.utils.telemetry_utils import log_streamlit_app_validation_retry, log_streamlit_app_conversion_success
 from mito_ai.path_utils import AbsoluteNotebookPath, AppFileName, get_absolute_notebook_dir_path, get_absolute_app_path, get_app_file_name
 
-async def generate_new_streamlit_code(notebook: List[dict]) -> str:
+async def generate_new_streamlit_code(notebook: List[dict], streamlit_app_prompt: str) -> str:
     """Send a query to the agent, get its response and parse the code"""
     
-    prompt_text = get_streamlit_app_creation_prompt(notebook)
+    prompt_text = get_streamlit_app_creation_prompt(notebook, streamlit_app_prompt)
     
     messages: List[MessageParam] = [
         cast(MessageParam, {
@@ -100,7 +100,7 @@ async def correct_error_in_generation(error: str, streamlit_app_code: str) -> st
 
     return streamlit_app_code
 
-async def streamlit_handler(notebook_path: AbsoluteNotebookPath, app_file_name: AppFileName, edit_prompt: str = "") -> None:
+async def streamlit_handler(create_new_app: bool, notebook_path: AbsoluteNotebookPath, app_file_name: AppFileName, streamlit_app_prompt: str = "") -> None:
     """Handler function for streamlit code generation and validation"""
 
     # Convert to absolute path for consistent handling
@@ -108,22 +108,22 @@ async def streamlit_handler(notebook_path: AbsoluteNotebookPath, app_file_name:
     app_directory = get_absolute_notebook_dir_path(notebook_path)
     app_path = get_absolute_app_path(app_directory, app_file_name)
     
-    if edit_prompt != "":
+    if create_new_app:
+        # Otherwise generate a new streamlit app
+        streamlit_code = await generate_new_streamlit_code(notebook_code, streamlit_app_prompt)
+    else:
         # If the user is editing an existing streamlit app, use the update function
-        streamlit_code = get_app_code_from_file(app_path)
+        existing_streamlit_code = get_app_code_from_file(app_path)
         
-        if streamlit_code is None:
+        if existing_streamlit_code is None:
             raise StreamlitConversionError("Error updating existing streamlit app because app.py file was not found.", 404)
         
-        streamlit_code = await update_existing_streamlit_code(notebook_code, streamlit_code, edit_prompt)
-    else:
-        # Otherwise generate a new streamlit app
-        streamlit_code = await generate_new_streamlit_code(notebook_code)
+        streamlit_code = await update_existing_streamlit_code(notebook_code, existing_streamlit_code, streamlit_app_prompt) 
        
     # Then, after creating/updating the app, validate that the new code runs 
     errors = validate_app(streamlit_code, notebook_path)
     tries = 0
-    while len(errors)>0 and tries < 5:
+    while len(errors) > 0 and tries < 5:
         for error in errors:
             streamlit_code = await correct_error_in_generation(error, streamlit_code)
         
@@ -141,4 +141,4 @@ async def streamlit_handler(notebook_path: AbsoluteNotebookPath, app_file_name:
     
     # Finally, update the app.py file with the new code
     create_app_file(app_path, streamlit_code)
-    log_streamlit_app_conversion_success('mito_server_key', MessageType.STREAMLIT_CONVERSION, edit_prompt)
+    log_streamlit_app_conversion_success('mito_server_key', MessageType.STREAMLIT_CONVERSION, streamlit_app_prompt)

mito_ai/streamlit_preview/handlers.py

@@ -22,12 +22,14 @@ class StreamlitPreviewHandler(APIHandler):
         self.preview_manager = StreamlitPreviewManager()
 
     @tornado.web.authenticated
+    
     async def post(self) -> None:
         """Start a new streamlit preview."""
         try:
+            
             # Parse and validate request
             body = self.get_json_body()
-            notebook_path, notebook_id, force_recreate, edit_prompt = validate_request_body(body)
+            notebook_path, notebook_id, force_recreate, streamlit_app_prompt = validate_request_body(body)
 
             # Ensure app exists
             absolute_notebook_path = get_absolute_notebook_path(notebook_path)
@@ -35,14 +37,19 @@ class StreamlitPreviewHandler(APIHandler):
             app_file_name = get_app_file_name(notebook_id)
             absolute_app_path = get_absolute_app_path(absolute_notebook_dir_path, app_file_name)
             app_path_exists = does_app_path_exist(absolute_app_path)
-
+            
             if not app_path_exists or force_recreate:
                 if not app_path_exists:
                     print("[Mito AI] App path not found, generating streamlit code")
                 else:
                     print("[Mito AI] Force recreating streamlit app")
 
-                await streamlit_handler(absolute_notebook_path, app_file_name, edit_prompt)
+                # Create a new app 
+                await streamlit_handler(True, absolute_notebook_path, app_file_name, streamlit_app_prompt)
+            elif streamlit_app_prompt != '':
+                # Update an existing app if there is a prompt provided. Otherwise, the user is just
+                # starting an existing app so we can skip the streamlit_handler all together
+                await streamlit_handler(False, absolute_notebook_path, app_file_name, streamlit_app_prompt)
 
             # Start preview
             # TODO: There's a bug here where when the user rebuilds and already running app. Instead of 
@@ -58,7 +65,7 @@ class StreamlitPreviewHandler(APIHandler):
                 "port": port, 
                 "url": f"http://localhost:{port}"
             })
-            log_streamlit_app_preview_success('mito_server_key', MessageType.STREAMLIT_CONVERSION, edit_prompt)
+            log_streamlit_app_preview_success('mito_server_key', MessageType.STREAMLIT_CONVERSION, streamlit_app_prompt)
 
         except StreamlitConversionError as e:
             print(e)
@@ -71,7 +78,7 @@ class StreamlitPreviewHandler(APIHandler):
                 MessageType.STREAMLIT_CONVERSION, 
                 error_message, 
                 formatted_traceback,
-                edit_prompt,
+                streamlit_app_prompt,
             )
         except StreamlitPreviewError as e:
             print(e)
@@ -79,7 +86,7 @@ class StreamlitPreviewHandler(APIHandler):
             formatted_traceback = traceback.format_exc()
             self.set_status(e.error_code)
             self.finish({"error": error_message})
-            log_streamlit_app_preview_failure('mito_server_key', MessageType.STREAMLIT_CONVERSION, error_message, formatted_traceback, edit_prompt)
+            log_streamlit_app_preview_failure('mito_server_key', MessageType.STREAMLIT_CONVERSION, error_message, formatted_traceback, streamlit_app_prompt)
         except Exception as e:
             print(f"Exception in streamlit preview handler: {e}")
             self.set_status(500)

mito_ai/streamlit_preview/manager.py

@@ -3,6 +3,7 @@
 
 import socket
 import subprocess
+import sys
 import time
 import threading
 import requests
@@ -54,8 +55,10 @@ class StreamlitPreviewManager:
             port = self.get_free_port()
             
             # Start streamlit process
+            # Use sys.executable -m streamlit to ensure it works on Windows
+            # where streamlit may not be directly executable in PATH
             cmd = [
-                "streamlit", "run", app_file_name,
+                sys.executable, "-m", "streamlit", "run", app_file_name,
                 "--server.port", str(port),
                 "--server.headless", "true",
                 "--server.address", "localhost",

mito_ai/streamlit_preview/utils.py

@@ -22,8 +22,8 @@ def validate_request_body(body: Optional[dict]) -> Tuple[str, str, bool, str]:
     if not isinstance(force_recreate, bool):
         raise StreamlitPreviewError("force_recreate must be a boolean", 400)
 
-    edit_prompt = body.get("edit_prompt", "")
-    if not isinstance(edit_prompt, str):
-        raise StreamlitPreviewError("edit_prompt must be a string", 400)
+    streamlit_app_prompt = body.get("streamlit_app_prompt", "")
+    if not isinstance(streamlit_app_prompt, str):
+        raise StreamlitPreviewError("streamlit_app_prompt must be a string", 400)
 
-    return notebook_path, notebook_id, force_recreate, edit_prompt
+    return notebook_path, notebook_id, force_recreate, streamlit_app_prompt

mito_ai/tests/streamlit_conversion/test_streamlit_agent_handler.py

@@ -89,7 +89,7 @@ class TestGenerateStreamlitCode:
         mock_stream.return_value = mock_async_gen()
         
         notebook_data: List[dict] = [{"cells": []}]
-        result = await generate_new_streamlit_code(notebook_data)
+        result = await generate_new_streamlit_code(notebook_data, '')
         
         expected_code = "import streamlit\nst.title('Hello')\n"
         assert result == expected_code
@@ -158,11 +158,11 @@ class TestStreamlitHandler:
         # Construct the expected app path using the same method as the production code
         app_directory = get_absolute_notebook_dir_path(notebook_path)
         expected_app_path = get_absolute_app_path(app_directory, app_file_name)
-        await streamlit_handler(notebook_path, app_file_name)
+        await streamlit_handler(True, notebook_path, app_file_name)
         
         # Verify calls
         mock_parse.assert_called_once_with(notebook_path)
-        mock_generate_code.assert_called_once_with(mock_notebook_data)
+        mock_generate_code.assert_called_once_with(mock_notebook_data, '')
         mock_validator.assert_called_once_with("import streamlit\nst.title('Test')", notebook_path)
         mock_create_file.assert_called_once_with(expected_app_path, "import streamlit\nst.title('Test')")
 
@@ -187,7 +187,7 @@ class TestStreamlitHandler:
     
         # Now it should raise an exception instead of returning a tuple
         with pytest.raises(Exception):
-            await streamlit_handler(AbsoluteNotebookPath("notebook.ipynb"), AppFileName('test-app-file-name.py'))
+            await streamlit_handler(True, AbsoluteNotebookPath("notebook.ipynb"), AppFileName('test-app-file-name.py'), '')
         
         # Verify that error correction was called 5 times (once per error, 5 retries)
         # Each retry processes 1 error, so 5 retries = 5 calls
@@ -215,7 +215,7 @@ class TestStreamlitHandler:
         
         # Now it should raise an exception instead of returning a tuple
         with pytest.raises(Exception):
-            await streamlit_handler(AbsoluteNotebookPath("notebook.ipynb"), AppFileName('test-app-file-name.py'))
+            await streamlit_handler(True, AbsoluteNotebookPath("notebook.ipynb"), AppFileName('test-app-file-name.py'), '')
 
     @pytest.mark.asyncio
     @patch('mito_ai.streamlit_conversion.streamlit_agent_handler.parse_jupyter_notebook_to_extract_required_content')
@@ -225,7 +225,7 @@ class TestStreamlitHandler:
         mock_parse.side_effect = FileNotFoundError("Notebook not found")
         
         with pytest.raises(FileNotFoundError, match="Notebook not found"):
-            await streamlit_handler(AbsoluteNotebookPath("notebook.ipynb"), AppFileName('test-app-file-name.py'))
+            await streamlit_handler(True, AbsoluteNotebookPath("notebook.ipynb"), AppFileName('test-app-file-name.py'), '')
 
     @pytest.mark.asyncio
     @patch('mito_ai.streamlit_conversion.streamlit_agent_handler.parse_jupyter_notebook_to_extract_required_content')
@@ -240,7 +240,7 @@ class TestStreamlitHandler:
         mock_generate_code.side_effect = Exception("Generation failed")
         
         with pytest.raises(Exception, match="Generation failed"):
-            await streamlit_handler(AbsoluteNotebookPath("notebook.ipynb"), AppFileName('test-app-file-name.py'))
+            await streamlit_handler(True, AbsoluteNotebookPath("notebook.ipynb"), AppFileName('test-app-file-name.py'), '')
 
 
 

mito_ai/tests/streamlit_preview/test_streamlit_preview_handler.py

@@ -99,9 +99,10 @@ class TestStreamlitPreviewHandler:
                     assert mock_streamlit_handler.called
                     # Verify it was called with the correct arguments
                     call_args = mock_streamlit_handler.call_args
-                    assert call_args[0][0] == os.path.abspath(notebook_path)  # First argument should be the absolute notebook path
-                    assert call_args[0][1] == app_file_name  # Second argument should be the app file name
-                    assert call_args[0][2] == ""  # Third argument should be the edit_prompt
+                    assert call_args[0][0] == True
+                    assert call_args[0][1] == os.path.abspath(notebook_path)  # First argument should be the absolute notebook path
+                    assert call_args[0][2] == app_file_name  # Second argument should be the app file name
+                    assert call_args[0][3] == ""  # Third argument should be the edit_prompt
                 else:
                     mock_streamlit_handler.assert_not_called()
                 

mito_ai/utils/open_ai_utils.py

@@ -171,6 +171,9 @@ def get_open_ai_completion_function_params(
         "stream": stream,
         "messages": messages,
     }
+
+    if model == "gpt-5.2":
+        completion_function_params["reasoning_effort"] = "low"
     
     # If a response format is provided, we need to convert it to a json schema.
     # Pydantic models are supported by the OpenAI API, however, we need to be able to