mito-ai 0.1.46__py3-none-any.whl → 0.1.49__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.46'
+__version__ = VERSION = '0.1.49'

mito_ai/app_deploy/app_deploy_utils.py

@@ -6,20 +6,39 @@ import zipfile
 import logging
 from typing import List, Optional
 
-def add_files_to_zip(zip_path: str, base_path: str, files_to_add: List[str], logger: Optional[logging.Logger] = None) -> None:
+from mito_ai.path_utils import AbsoluteNotebookDirPath
+
+def add_files_to_zip(
+    zip_path: str, 
+    notebook_dir_path: AbsoluteNotebookDirPath,
+    files_to_add: List[str], 
+    app_file_name: str,
+    logger: Optional[logging.Logger] = None
+) -> None:
     """Create a zip file at zip_path and add the selected files/folders."""
     with zipfile.ZipFile(zip_path, "w", zipfile.ZIP_DEFLATED) as zipf:
-        for rel_path in files_to_add:
-            abs_path = os.path.join(base_path, rel_path)
+        for file_to_add_rel_path in files_to_add:
+            
+            file_to_add_abs_path = os.path.join(notebook_dir_path, file_to_add_rel_path)
 
-            if os.path.isfile(abs_path):
-                zipf.write(abs_path, arcname=rel_path)
-            elif os.path.isdir(abs_path):
-                for root, _, files in os.walk(abs_path):
+            if os.path.isfile(file_to_add_abs_path):
+                basename = os.path.basename(file_to_add_abs_path)
+                
+                if basename == app_file_name:
+                    # For the actual app file, we want to write it just as app.py 
+                    # so our infra can always deploy using `streamlit run app.py`
+                    # without having to account for different app names
+                    zipf.write(file_to_add_abs_path, arcname='app.py')
+                else:
+                    # otherwise we want to keep the name as is so all references
+                    # to it from the app are correct
+                    zipf.write(file_to_add_abs_path, arcname=file_to_add_rel_path)
+            elif os.path.isdir(file_to_add_abs_path):
+                for root, _, files in os.walk(file_to_add_abs_path):
                     for file in files:
                         file_abs = os.path.join(root, file)
-                        arcname = os.path.relpath(file_abs, base_path)
+                        arcname = os.path.relpath(file_abs, notebook_dir_path)
                         zipf.write(file_abs, arcname=arcname)
             else:
                 if logger:
-                    logger.warning(f"Skipping missing file: {abs_path}")
+                    logger.warning(f"Skipping missing file: {file_to_add_abs_path}")

mito_ai/app_deploy/handlers.py

@@ -4,10 +4,11 @@
 import os
 import time
 import logging
-from typing import Any, Union, List
+from typing import Any, Union, List, Optional
 import tempfile
-from mito_ai.streamlit_conversion.streamlit_utils import get_app_path
+from mito_ai.path_utils import AbsoluteNotebookDirPath, AppFileName, does_app_path_exist, get_absolute_app_path, get_absolute_notebook_dir_path, get_absolute_notebook_path, get_app_file_name
 from mito_ai.utils.create import initialize_user
+from mito_ai.utils.error_classes import StreamlitDeploymentError
 from mito_ai.utils.version_utils import is_pro
 from mito_ai.utils.websocket_base import BaseWebSocketHandler
 from mito_ai.app_deploy.app_deploy_utils import  add_files_to_zip
@@ -16,11 +17,13 @@ from mito_ai.app_deploy.models import (
     AppDeployError,
     DeployAppRequest,
     ErrorMessage,
-    MessageType
 )
+from mito_ai.completions.models import MessageType
 from mito_ai.logger import get_logger
 from mito_ai.constants import ACTIVE_STREAMLIT_BASE_URL
+from mito_ai.utils.telemetry_utils import log_streamlit_app_deployment_failure
 import requests
+import traceback
 
 
 class AppDeployHandler(BaseWebSocketHandler):
@@ -67,41 +70,46 @@ class AppDeployHandler(BaseWebSocketHandler):
         self.log.debug("App builder message received: %s", message)
         
         try:
-            # Ensure message is a string before parsing
-            if not isinstance(message, str):
-                raise ValueError("Message must be a string")
-
             parsed_message = self.parse_message(message)
             message_type = parsed_message.get('type')
             
             if message_type == MessageType.DEPLOY_APP.value:
                 # Handle build app request
                 deploy_app_request = DeployAppRequest(**parsed_message)
-                await self._handle_deploy_app(deploy_app_request)
+                response = await self._handle_deploy_app(deploy_app_request)
+                self.reply(response)
             else:
                 self.log.error(f"Unknown message type: {message_type}")
                 error = AppDeployError(
                     error_type="InvalidRequest",
-                    title=f"Unknown message type: {message_type}"
+                    message=f"Unknown message type: {message_type}",
+                    error_code=400
                 )
-                self.reply(ErrorMessage(**error.__dict__))
-                
-        except ValueError as e:
+                raise StreamlitDeploymentError(error)
+
+        except StreamlitDeploymentError as e:
             self.log.error("Invalid app builder request", exc_info=e)
-            error = AppDeployError.from_exception(e)
-            self.reply(ErrorMessage(**error.__dict__))
+            log_streamlit_app_deployment_failure('mito_server_key', MessageType.DEPLOY_APP, e.error.__dict__)
+            self.reply(
+                DeployAppReply(
+                        parent_id=e.message_id,
+                        url="",
+                        error=ErrorMessage(**e.error.__dict__)
+                    )
+            )
         except Exception as e:
             self.log.error("Error handling app builder message", exc_info=e)
             error = AppDeployError.from_exception(
                 e, 
                 hint="An error occurred while building the app. Please check the logs for details."
             )
+            log_streamlit_app_deployment_failure('mito_server_key', MessageType.DEPLOY_APP, error.__dict__)
             self.reply(ErrorMessage(**error.__dict__))
             
         latency_ms = round((time.time() - start) * 1000)
         self.log.info(f"App builder handler processed in {latency_ms} ms.")
     
-    async def _handle_deploy_app(self, message: DeployAppRequest) -> None:
+    async def _handle_deploy_app(self, message: DeployAppRequest) -> DeployAppReply:
         """Handle a build app request.
         
         Args:
@@ -109,24 +117,29 @@ class AppDeployHandler(BaseWebSocketHandler):
         """
         message_id = message.message_id
         notebook_path = message.notebook_path
+        notebook_id = message.notebook_id
         jwt_token = message.jwt_token
         files_to_upload = message.selected_files
         
+        # Validate parameters
+        missing_required_parameters = []
         if not message_id:
-            self.log.error("Missing message_id in request")
-            return
-        
+            missing_required_parameters.append('message_id')
+        if not notebook_id:
+            missing_required_parameters.append('notebook_id')
         if not notebook_path:
+            missing_required_parameters.append('notebook_path')
+            
+        if len(missing_required_parameters) > 0:
+            error_message = f"Missing required request parameters: {', '.join(missing_required_parameters)}"
+            self.log.error(error_message)
             error = AppDeployError(
-                error_type="InvalidRequest",
-                title="Missing 'notebook_path' parameter"
+                error_type="BadRequest",
+                message=error_message,
+                error_code=400,
+                message_id=message_id
             )
-            self.reply(DeployAppReply(
-                parent_id=message_id,
-                url="",
-                error=error
-            ))
-            return
+            raise StreamlitDeploymentError(error)
 
         # Validate JWT token if provided
         token_preview = jwt_token[:20] if jwt_token else "No token provided"
@@ -136,55 +149,49 @@ class AppDeployHandler(BaseWebSocketHandler):
             self.log.error("JWT token validation failed")
             error = AppDeployError(
                 error_type="Unauthorized",
-                title="Invalid authentication token",
-                hint="Please sign in again to deploy your app."
+                message="Invalid authentication token",
+                hint="Please sign in again to deploy your app.",
+                error_code=401,
+                message_id=message_id
             )
-            self.reply(DeployAppReply(
-                parent_id=message_id,
-                url="",
-                error=error
-            ))
-            return
+            raise StreamlitDeploymentError(error)
         else:
             self.log.info("JWT token validation successful")
-            
-        try:
-            notebook_path = str(notebook_path) if notebook_path else ""
 
-            app_directory = os.path.dirname(notebook_path)
-            
-            # Check if the app.py file exists
-            app_path = get_app_path(app_directory)
-            if app_path is None:
-                error = AppDeployError(
-                    error_type="AppNotFound",
-                    title="App not found",
-                    hint="Please make sure the app.py file exists in the same directory as the notebook."
-                )
-                self.reply(DeployAppReply(
-                    parent_id=message_id,
-                    url="",
-                    error=error
-                ))
-            
-            # Finally, deploy the app
-            deploy_url = await self._deploy_app(app_directory, files_to_upload, jwt_token)
+        notebook_path = str(notebook_path) if notebook_path else ""
+        absolute_notebook_path = get_absolute_notebook_path(notebook_path)
+        absolute_app_directory = get_absolute_notebook_dir_path(absolute_notebook_path)
+        app_file_name = get_app_file_name(notebook_id)
+        app_path = get_absolute_app_path(absolute_app_directory, app_file_name)
+
+        # Check if the app.py file exists
+        app_path_exists = does_app_path_exist(app_path)
+        if not app_path_exists:
+            error = AppDeployError(
+                error_type="AppNotFound",
+                message="App not found",
+                hint=f"Please make sure the {app_file_name} file exists in the same directory as the notebook.",
+                error_code=400,
+                message_id=message_id
+            )
+            raise StreamlitDeploymentError(error)
+
+        # Finally, deploy the app
+        deploy_url = await self._deploy_app(
+            absolute_app_directory, 
+            app_file_name,
+            files_to_upload, 
+            message_id,
+            jwt_token
+        )
+
+        # Send the response
+        return DeployAppReply(
+            parent_id=message_id,
+            url=deploy_url if deploy_url else ""
+        )
+        
 
-            # Send the response
-            self.reply(DeployAppReply(
-                parent_id=message_id,
-                url=deploy_url
-            ))
-            
-        except Exception as e:
-            self.log.error(f"Error building app: {e}", exc_info=e)
-            error = AppDeployError.from_exception(e)
-            self.reply(DeployAppReply(
-                parent_id=message_id,
-                url="",
-                error=error
-            ))
-    
     def _validate_jwt_token(self, token: str) -> bool:
         """Basic JWT token validation logic.
 
@@ -219,7 +226,15 @@ class AppDeployHandler(BaseWebSocketHandler):
             return False
 
 
-    async def _deploy_app(self, app_path: str, files_to_upload:List[str], jwt_token: str = '') -> str:
+    async def _deploy_app(
+        self,
+        absolute_notebook_dir_path: AbsoluteNotebookDirPath, 
+        app_file_name: AppFileName, 
+        files_to_upload:List[str], 
+        message_id: str, 
+        jwt_token: str = ''
+    ) -> Optional[str]:
+        
         """Deploy the app using pre-signed URLs.
         
         Args:
@@ -230,9 +245,11 @@ class AppDeployHandler(BaseWebSocketHandler):
         Returns:
             The URL of the deployed app.
         """
-        # Get app name from the path
-        app_name = os.path.basename(app_path).split('.')[0]
-        self.log.info(f"Deploying app: {app_name} from path: {app_path}")
+        # Get app name from the path without the file type ending 
+        # ie: if the file is my-app.py, this variable is just my-app
+        # We use it in the app url
+        app_file_name_no_file_extension_ending = app_file_name.split('.')[0]
+        self.log.info(f"Deploying app: {app_file_name} from path: {absolute_notebook_dir_path}")
         
         try:
             # Step 1: Get pre-signed URL from API
@@ -247,7 +264,7 @@ class AppDeployHandler(BaseWebSocketHandler):
 
             headers["Subscription-Tier"] = 'Pro' if is_pro() else 'Standard'
 
-            url_response = requests.get(f"{ACTIVE_STREAMLIT_BASE_URL}/get-upload-url?app_name={app_name}", headers=headers)
+            url_response = requests.get(f"{ACTIVE_STREAMLIT_BASE_URL}/get-upload-url?app_name={app_file_name_no_file_extension_ending}", headers=headers)
             url_response.raise_for_status()
             
             url_data = url_response.json()
@@ -255,7 +272,7 @@ class AppDeployHandler(BaseWebSocketHandler):
             expected_app_url = url_data['expected_app_url']
             
             self.log.info(f"Received pre-signed URL. App will be available at: {expected_app_url}")
-            
+
             # Step 2: Create a zip file of the app.
             temp_zip_path = None
             try:
@@ -264,12 +281,19 @@ class AppDeployHandler(BaseWebSocketHandler):
                     temp_zip_path = temp_zip.name
 
                 self.log.info("Zipping application files...")
-                add_files_to_zip(temp_zip_path, app_path, files_to_upload, self.log)
+                add_files_to_zip(temp_zip_path, absolute_notebook_dir_path, files_to_upload, app_file_name, self.log)
 
                 upload_response = await self._upload_app_to_s3(temp_zip_path, presigned_url)
             except Exception as e:
                 self.log.error(f"Error zipping app: {e}")
-                raise
+                error = AppDeployError(
+                    error_type="ZippingError",
+                    message=f"Error zipping app: {e}",
+                    traceback=traceback.format_exc(),
+                    error_code=500,
+                    message_id=message_id
+                )
+                raise StreamlitDeploymentError(error)
             finally:
                 # Clean up
                 if temp_zip_path is not None:
@@ -278,20 +302,35 @@ class AppDeployHandler(BaseWebSocketHandler):
             self.log.info(f"Upload successful! Status code: {upload_response.status_code}")
             
             self.log.info(f"Deployment initiated. App will be available at: {expected_app_url}")
-            return expected_app_url # type: ignore
+            return str(expected_app_url)
             
         except requests.exceptions.RequestException as e:
             self.log.error(f"Error during API request: {e}")
             if hasattr(e, 'response') and e.response is not None:
                 error_detail = e.response.json()
+                error_message = error_detail.get('error', "")
                 self.log.error(f"Server error details: {error_detail}")
-                if 'error' in error_detail:
-                    raise Exception(error_detail['error'])
-                raise
+            else:
+                error_message = str(e)
+            
+            error = AppDeployError(
+                error_type="APIException",
+                message=str(error_message),
+                traceback=traceback.format_exc(),
+                error_code=500,
+                message_id=message_id
+            )
+            raise StreamlitDeploymentError(error)
         except Exception as e:
             self.log.error(f"Error during deployment: {str(e)}")
-            raise
-        raise RuntimeError("Unexpected error in _deploy_app")
+            error = AppDeployError(
+                error_type="DeploymentException",
+                message=str(e),
+                traceback=traceback.format_exc(),
+                error_code=500,
+                message_id=message_id
+            )
+            raise StreamlitDeploymentError(error)
 
     async def _upload_app_to_s3(self, app_path: str, presigned_url: str) -> requests.Response:
         """Upload the app to S3 using the presigned URL."""

mito_ai/app_deploy/models.py

@@ -2,15 +2,10 @@
 # Distributed under the terms of the GNU Affero General Public License v3.0 License.
 
 from dataclasses import dataclass
-from enum import Enum
+import traceback
 from typing import Literal, Optional, List
 
 
-class MessageType(str, Enum):
-    """Types of app deploy messages."""
-    DEPLOY_APP = "deploy-app"
-
-
 @dataclass(frozen=True)
 class AppDeployError:
     """Error information for app deploy operations."""
@@ -19,7 +14,13 @@ class AppDeployError:
     error_type: str
     
     # Error title.
-    title: str
+    message: str
+
+    #ID of parent to resolve response
+    message_id: Optional[str] = "InvalidMessageID"
+
+    # Error code
+    error_code: Optional[int] = 500
     
     # Error traceback.
     traceback: Optional[str] = None
@@ -28,21 +29,24 @@ class AppDeployError:
     hint: Optional[str] = None
     
     @classmethod
-    def from_exception(cls, e: Exception, hint: Optional[str] = None) -> "AppDeployError":
+    def from_exception(cls, e: Exception, hint: Optional[str] = None, error_code: Optional[int] = 500) -> "AppDeployError":
         """Create an error from an exception.
 
         Args:
             e: The exception.
             hint: Optional hint to fix the error.
+            error_code: Optional error code which defaults to 500
 
         Returns:
             The app builder error.
         """
+        tb_str = "".join(traceback.format_exception(type(e), e, e.__traceback__))
         return cls(
             error_type=type(e).__name__,
-            title=str(e),
-            traceback=getattr(e, "__traceback__", None) and str(e.__traceback__),
+            message=str(e),
+            traceback=tb_str,
             hint=hint,
+            error_code=error_code
         )
 
 
@@ -59,13 +63,16 @@ class DeployAppRequest:
     """Request to deploy an app."""
     
     # Request type.
-    type: Literal["deploy-app"]
+    type: Literal["deploy_app"]
     
     # Message ID.
     message_id: str
     
     # Path to the app file.
     notebook_path: str
+    
+    # Notebook ID
+    notebook_id: str
 
     # Files to be uploaded for the app to run
     selected_files: List[str]
@@ -88,4 +95,4 @@ class DeployAppReply:
     error: Optional[AppDeployError] = None
     
     # Type of reply.
-    type: Literal["deploy-app"] = "deploy-app" 
+    type: Literal["deploy_app"] = "deploy_app"

mito_ai/completions/models.py

@@ -29,12 +29,13 @@ class CellUpdate(BaseModel):
 # for now and rely on the AI to respond with the correct types, following the format
 # that we show it in the system prompt.
 class AgentResponse(BaseModel):
-    type: Literal['cell_update', 'get_cell_output', 'run_all_cells', 'finished_task']
+    type: Literal['cell_update', 'get_cell_output', 'run_all_cells', 'finished_task', 'create_streamlit_app', 'edit_streamlit_app']
     message: str
     cell_update: Optional[CellUpdate]
     get_cell_output_cell_id: Optional[str]
     next_steps: Optional[List[str]]
     analysis_assumptions: Optional[List[str]]
+    edit_streamlit_app_prompt: Optional[str]
     
     
 @dataclass(frozen=True)
@@ -65,6 +66,7 @@ class MessageType(Enum):
     UPDATE_MODEL_CONFIG = "update_model_config"
     STREAMLIT_CONVERSION = "streamlit_conversion"
     STOP_AGENT = "stop_agent"
+    DEPLOY_APP = "deploy_app"
 
     
 @dataclass(frozen=True)
@@ -95,7 +97,10 @@ class AgentExecutionMetadata():
     threadId: ThreadID
     input: str
     aiOptimizedCells: List[AIOptimizedCell]
+    activeCellId: str
     isChromeBrowser: bool
+    notebookPath: str
+    notebookID: str
     base64EncodedActiveCellOutput: Optional[str] = None
     variables: Optional[List[str]] = None
     files: Optional[List[str]] = None

mito_ai/completions/prompt_builders/agent_execution_prompt.py

@@ -3,14 +3,17 @@
 
 from mito_ai.completions.models import AgentExecutionMetadata
 from mito_ai.completions.prompt_builders.prompt_constants import (
+    ACTIVE_CELL_ID_SECTION_HEADING,
     FILES_SECTION_HEADING,
     JUPYTER_NOTEBOOK_SECTION_HEADING,
+    STREAMLIT_APP_STATUS_SECTION_HEADING,
     VARIABLES_SECTION_HEADING,
     cell_update_output_str
 )
 from mito_ai.completions.prompt_builders.utils import (
     get_rules_str,
     get_selected_context_str,
+    get_streamlit_app_status_str
 )
 
 
@@ -20,7 +23,10 @@ def create_agent_execution_prompt(md: AgentExecutionMetadata) -> str:
     ai_optimized_cells_str = '\n'.join([f"{cell}" for cell in md.aiOptimizedCells or []])
     rules_str = get_rules_str(md.additionalContext)
     selected_context_str = get_selected_context_str(md.additionalContext)
-        
+    
+
+    streamlit_status_str = get_streamlit_app_status_str(md.notebookID, md.notebookPath)    
+    
     context_str = f"""Remember to choose the correct tool to respond with.
 
 {rules_str}
@@ -35,6 +41,12 @@ def create_agent_execution_prompt(md: AgentExecutionMetadata) -> str:
 {FILES_SECTION_HEADING}
 {files_str}
 
+{STREAMLIT_APP_STATUS_SECTION_HEADING}
+{streamlit_status_str}
+
+{ACTIVE_CELL_ID_SECTION_HEADING}
+{md.activeCellId}
+
 {selected_context_str}
 
 {cell_update_output_str(md.base64EncodedActiveCellOutput is not None)}"""

mito_ai/completions/prompt_builders/agent_system_message.py

@@ -225,6 +225,57 @@ Important information:
 5. Do not use this tool repeatedly if it continues to produce errors - instead, focus on fixing the specific error that occurred.
 ====
 
+TOOL: CREATE_STREAMLIT_APP
+
+When you want to create a new Streamlit app from the current notebook, respond with this format:
+
+{{
+    type: 'create_streamlit_app',
+    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.
+
+<Example>
+
+Your task: Show me my notebook as an app.
+
+Output:
+{{
+    type: 'create_streamlit_app',
+    message: "I'll convert your notebook into an app."
+}}
+
+The user will see a preview of the app and because you fulfilled your task, you can next respond with a FINISHED_TASK tool message.
+
+<Example>
+
+====
+
+TOOL: EDIT_STREAMLIT_APP
+
+When you want to edit an existing Streamlit app, respond with this format:
+
+{{
+    type: 'edit_streamlit_app',
+    message: str,
+    edit_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").
+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. 
+
+====
+
 TOOL: FINISHED_TASK
 
 When you have completed the user's task, respond with a message in this format:
@@ -238,8 +289,8 @@ 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.
-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: ["Visualize the results", "Export the cleaned data to CSV", "Perform statistical analysis on the key metrics"].
-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 customer age group"].
+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.
 6. If the user's task doesn't involve creating or modifying a code cell, you should respond with a FINISHED_TASK response. 
 7. If the user is just sending a friendly greeting (like "Hello", "Hi", "Hey", "How are you?", "What can you help me with?", etc.), you must respond with a FINISHED_TASK response message with a friendly message like this: "Hello! I'm Mito AI, your AI assistant for data analysis and Python programming in Jupyter notebooks. I can help you analyze datasets, create visualizations, clean data, and much more. What would you like to work on today?"
@@ -391,7 +442,9 @@ As you are guiding the user through the process of completing the task, send the
 
 The user is a beginning Python user, so you will need to be careful to send them only small steps to complete. Don't try to complete the task in a single response to the user. Instead, each message you send to the user should only contain a single, small step towards the end goal. When the user has completed the step, they will let you know that they are ready for the next step. 
 
-You will keep working in the following iterative format until you have decided that you have finished the user's request. When you decide that you have finished the user's request, respond with a FINISHED_TASK tool message. Otherwise, if you have not finished the user's request, respond with a CELL_UPDATE {OR_GET_CELL_OUTPUT} tool message. When you respond with a CELL_UPDATE, the user will apply the CELL_UPDATE to the notebook and run the new code cell. The user will then send you a message with an updated version of the variables defined in the kernel, code in the notebook, and files in the current directory. In addition, the user will check if the code you provided produced an errored when executed. If it did produce an error, the user will share the error message with you.
+You will keep working in the following iterative format until you have decided that you have finished the user's request. When you decide that you have finished the user's request, respond with a FINISHED_TASK tool message. Otherwise, if you have not finished the user's request, respond with one of your other tools. 
+
+When you respond with a CELL_UPDATE, the user will apply the CELL_UPDATE to the notebook and run the new code cell. The user will then send you a message with an updated version of the variables defined in the kernel, code in the notebook, and files in the current directory. In addition, the user will check if the code you provided produced an errored when executed. If it did produce an error, the user will share the error message with you.
 
 Whenever you get a message back from the user, you should:
 1. Ask yourself if the previous message you sent to the user was correct. You can answer this question by reviewing the updated code, variables, or output of the cell if you requested it.
@@ -410,4 +463,10 @@ REMEMBER, YOU ARE GOING TO COMPLETE THE USER'S TASK OVER THE COURSE OF THE ENTIR
 - If you are happy with the analysis, refer back to the original task provided by the user to decide your next steps. In this example, it is to graph the results, so you will send a CellAddition to construct the graph. 
 - Wait for the user to send you back the updated variables and notebook state.
 {'' if not isChromeBrowser else '- Send a GET_CELL_OUTPUT tool message to get the output of the cell you just created and check if you can improve the graph to make it more readable, informative, or professional.'}
-- If after reviewing the updates you decide that you've completed the task, send a FINISHED_TASK tool message."""
+- If after reviewing the updates you decide that you've completed the task, send a FINISHED_TASK tool message.
+
+====
+
+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."""

mito_ai/completions/prompt_builders/chat_system_message.py

@@ -21,6 +21,10 @@ There are three possible types of responses you might give:
 2. Explanation/Analysis: If the task does not require a code update, it might instead require you to provide an explanation of existing code or data, provide an analysis of the the data or chart.
 3. Friendly Response: If the user is just asking a question, saying hi, or you're just chatting, respond with a friendly response and do not return any code.
 
+Other useful information:
+1. The user has two types of modes that they can collaborate with you in: Chat Mode (this mode) and agent mode. Chat mode gives the user more control over the edits made to the notebook and only edits the active cell. Agent mode gives you more autonomy over completing the user's task across mulitple messages. In agent mode, you can edit or create new cells, see the entire notebook, automatically run the code you write, and more.
+2. If the user asks you to generate a dashboard, app, or streamlit app for them, you should tell them that they must use Agent mode to complete the task. You are not able to automatically switch the user to agent mode, but they can switch to it themselves by using the Chat/Agent mode toggle in the bottom left corner of the Ai taskpane.
+
 ====
 {CITATION_RULES}