Merge pull request #2078 from mito-ds/dev

Release December 5, 2025

Author: ngafar

mito-ai/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/mito_ai/completions/prompt_builders/agent_system_message.py

@@ -231,15 +231,17 @@ def create_agent_system_message_prompt(isChromeBrowser: bool) -> str:
 
 {{
     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 +250,7 @@ def create_agent_system_message_prompt(isChromeBrowser: bool) -> str:
 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 +267,12 @@ def create_agent_system_message_prompt(isChromeBrowser: bool) -> str:
 {{
     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. 

mito-ai/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 @@ def get_streamlit_app_creation_prompt(notebook: List[dict]) -> str:
 ]
 </Example>
 
-Notebook to convert:
+{streamlit_app_spec_section}
+
+NOTEBOOK TO CONVERT:
 
 {notebook}
 """

mito-ai/mito_ai/streamlit_conversion/streamlit_agent_handler.py

@@ -16,10 +16,10 @@
 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,30 +100,30 @@ 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
     notebook_code = parse_jupyter_notebook_to_extract_required_content(notebook_path)
     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/mito_ai/streamlit_preview/handlers.py

@@ -22,27 +22,34 @@ def initialize(self) -> None:
         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)
             absolute_notebook_dir_path = get_absolute_notebook_dir_path(absolute_notebook_path)
             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 @@ async def post(self) -> None:
                 "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,15 +78,15 @@ async def post(self) -> None:
                 MessageType.STREAMLIT_CONVERSION, 
                 error_message, 
                 formatted_traceback,
-                edit_prompt,
+                streamlit_app_prompt,
             )
         except StreamlitPreviewError as e:
             print(e)
             error_message = str(e)
             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/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 @@ def start_streamlit_preview(self, app_directory: AbsoluteNotebookDirPath, app_fi
             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/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/mito_ai/tests/streamlit_conversion/test_streamlit_agent_handler.py

@@ -89,7 +89,7 @@ async def mock_async_gen():
         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 @@ async def test_streamlit_handler_success(self, mock_log_success, mock_create_fil
         # 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 @@ async def test_streamlit_handler_max_retries_exceeded(self, mock_log_retry, mock
 
         # 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 @@ async def test_streamlit_handler_file_creation_failure(self, mock_create_file, m
 
         # 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 @@ async def test_streamlit_handler_parse_notebook_exception(self, mock_parse):
         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 @@ async def test_streamlit_handler_generation_exception(self, mock_generate_code,
         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/mito_ai/tests/streamlit_preview/test_streamlit_preview_handler.py

@@ -99,9 +99,10 @@ def mock_finish_func(response):
                     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/package.json

@@ -34,8 +34,8 @@
     "build:prod": "jlpm clean && jlpm build:lib:prod && jlpm build:labextension",
     "build:labextension": "jupyter labextension build .",
     "build:labextension:dev": "jupyter labextension build --development True .",
-    "build:lib": "rm -rf buildcache && npx tsc --sourceMap",
-    "build:lib:prod": "rm -rf buildcache && npx tsc",
+    "build:lib": "rimraf buildcache && npx tsc --sourceMap",
+    "build:lib:prod": "rimraf buildcache && npx tsc",
     "clean": "jlpm clean:lib",
     "clean:lib": "rimraf lib tsconfig.tsbuildinfo",
     "clean:lintcache": "rimraf .eslintcache .stylelintcache",