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

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.

Potentially problematic release.


This version of mito-ai might be problematic. Click here for more details.

Files changed (70) hide show
  1. mito_ai/_version.py +1 -1
  2. mito_ai/app_deploy/handlers.py +97 -77
  3. mito_ai/app_deploy/models.py +16 -12
  4. mito_ai/completions/models.py +4 -1
  5. mito_ai/completions/prompt_builders/agent_execution_prompt.py +6 -1
  6. mito_ai/completions/prompt_builders/agent_system_message.py +63 -4
  7. mito_ai/completions/prompt_builders/chat_system_message.py +4 -0
  8. mito_ai/completions/prompt_builders/prompt_constants.py +1 -0
  9. mito_ai/completions/prompt_builders/utils.py +14 -0
  10. mito_ai/path_utils.py +56 -0
  11. mito_ai/streamlit_conversion/agent_utils.py +4 -201
  12. mito_ai/streamlit_conversion/prompts/prompt_constants.py +142 -152
  13. mito_ai/streamlit_conversion/prompts/streamlit_error_correction_prompt.py +3 -3
  14. mito_ai/streamlit_conversion/prompts/streamlit_finish_todo_prompt.py +2 -2
  15. mito_ai/streamlit_conversion/prompts/update_existing_app_prompt.py +2 -2
  16. mito_ai/streamlit_conversion/search_replace_utils.py +93 -0
  17. mito_ai/streamlit_conversion/streamlit_agent_handler.py +29 -39
  18. mito_ai/streamlit_conversion/streamlit_utils.py +11 -64
  19. mito_ai/streamlit_conversion/validate_streamlit_app.py +5 -18
  20. mito_ai/streamlit_preview/handlers.py +44 -84
  21. mito_ai/streamlit_preview/manager.py +6 -6
  22. mito_ai/streamlit_preview/utils.py +16 -19
  23. mito_ai/tests/streamlit_conversion/test_apply_search_replace.py +226 -0
  24. mito_ai/tests/streamlit_conversion/test_streamlit_agent_handler.py +29 -53
  25. mito_ai/tests/streamlit_conversion/test_streamlit_utils.py +26 -29
  26. mito_ai/tests/streamlit_conversion/test_validate_streamlit_app.py +6 -3
  27. mito_ai/tests/streamlit_preview/test_streamlit_preview_handler.py +12 -15
  28. mito_ai/tests/streamlit_preview/test_streamlit_preview_manager.py +22 -26
  29. mito_ai/user/handlers.py +15 -3
  30. mito_ai/utils/create.py +17 -1
  31. mito_ai/utils/error_classes.py +42 -0
  32. mito_ai/utils/message_history_utils.py +3 -1
  33. mito_ai/utils/telemetry_utils.py +75 -10
  34. {mito_ai-0.1.46.data → mito_ai-0.1.47.data}/data/share/jupyter/labextensions/mito_ai/build_log.json +1 -1
  35. {mito_ai-0.1.46.data → mito_ai-0.1.47.data}/data/share/jupyter/labextensions/mito_ai/package.json +2 -2
  36. {mito_ai-0.1.46.data → mito_ai-0.1.47.data}/data/share/jupyter/labextensions/mito_ai/schemas/mito_ai/package.json.orig +1 -1
  37. mito_ai-0.1.46.data/data/share/jupyter/labextensions/mito_ai/static/lib_index_js.20f12766ecd3d430568e.js → mito_ai-0.1.47.data/data/share/jupyter/labextensions/mito_ai/static/lib_index_js.2db61d2b629817845901.js +1274 -293
  38. mito_ai-0.1.47.data/data/share/jupyter/labextensions/mito_ai/static/lib_index_js.2db61d2b629817845901.js.map +1 -0
  39. mito_ai-0.1.46.data/data/share/jupyter/labextensions/mito_ai/static/remoteEntry.54126ab6511271265443.js → mito_ai-0.1.47.data/data/share/jupyter/labextensions/mito_ai/static/remoteEntry.e22c6cd4e56c32116daa.js +7 -7
  40. mito_ai-0.1.46.data/data/share/jupyter/labextensions/mito_ai/static/remoteEntry.54126ab6511271265443.js.map → mito_ai-0.1.47.data/data/share/jupyter/labextensions/mito_ai/static/remoteEntry.e22c6cd4e56c32116daa.js.map +1 -1
  41. {mito_ai-0.1.46.dist-info → mito_ai-0.1.47.dist-info}/METADATA +1 -1
  42. {mito_ai-0.1.46.dist-info → mito_ai-0.1.47.dist-info}/RECORD +67 -65
  43. mito_ai/tests/streamlit_conversion/test_apply_patch_to_text.py +0 -368
  44. mito_ai/tests/streamlit_conversion/test_fix_diff_headers.py +0 -533
  45. mito_ai-0.1.46.data/data/share/jupyter/labextensions/mito_ai/static/lib_index_js.20f12766ecd3d430568e.js.map +0 -1
  46. /mito_ai/streamlit_conversion/{streamlit_system_prompt.py → prompts/streamlit_system_prompt.py} +0 -0
  47. {mito_ai-0.1.46.data → mito_ai-0.1.47.data}/data/etc/jupyter/jupyter_server_config.d/mito_ai.json +0 -0
  48. {mito_ai-0.1.46.data → mito_ai-0.1.47.data}/data/share/jupyter/labextensions/mito_ai/schemas/mito_ai/toolbar-buttons.json +0 -0
  49. {mito_ai-0.1.46.data → mito_ai-0.1.47.data}/data/share/jupyter/labextensions/mito_ai/static/node_modules_process_browser_js.4b128e94d31a81ebd209.js +0 -0
  50. {mito_ai-0.1.46.data → mito_ai-0.1.47.data}/data/share/jupyter/labextensions/mito_ai/static/node_modules_process_browser_js.4b128e94d31a81ebd209.js.map +0 -0
  51. {mito_ai-0.1.46.data → mito_ai-0.1.47.data}/data/share/jupyter/labextensions/mito_ai/static/style.js +0 -0
  52. {mito_ai-0.1.46.data → mito_ai-0.1.47.data}/data/share/jupyter/labextensions/mito_ai/static/style_index_js.5876024bb17dbd6a3ee6.js +0 -0
  53. {mito_ai-0.1.46.data → mito_ai-0.1.47.data}/data/share/jupyter/labextensions/mito_ai/static/style_index_js.5876024bb17dbd6a3ee6.js.map +0 -0
  54. {mito_ai-0.1.46.data → mito_ai-0.1.47.data}/data/share/jupyter/labextensions/mito_ai/static/vendors-node_modules_aws-amplify_auth_dist_esm_providers_cognito_apis_signOut_mjs-node_module-75790d.688c25857e7b81b1740f.js +0 -0
  55. {mito_ai-0.1.46.data → mito_ai-0.1.47.data}/data/share/jupyter/labextensions/mito_ai/static/vendors-node_modules_aws-amplify_auth_dist_esm_providers_cognito_apis_signOut_mjs-node_module-75790d.688c25857e7b81b1740f.js.map +0 -0
  56. {mito_ai-0.1.46.data → mito_ai-0.1.47.data}/data/share/jupyter/labextensions/mito_ai/static/vendors-node_modules_aws-amplify_auth_dist_esm_providers_cognito_tokenProvider_tokenProvider_-72f1c8.a917210f057fcfe224ad.js +0 -0
  57. {mito_ai-0.1.46.data → mito_ai-0.1.47.data}/data/share/jupyter/labextensions/mito_ai/static/vendors-node_modules_aws-amplify_auth_dist_esm_providers_cognito_tokenProvider_tokenProvider_-72f1c8.a917210f057fcfe224ad.js.map +0 -0
  58. {mito_ai-0.1.46.data → mito_ai-0.1.47.data}/data/share/jupyter/labextensions/mito_ai/static/vendors-node_modules_aws-amplify_dist_esm_index_mjs.6bac1a8c4cc93f15f6b7.js +0 -0
  59. {mito_ai-0.1.46.data → mito_ai-0.1.47.data}/data/share/jupyter/labextensions/mito_ai/static/vendors-node_modules_aws-amplify_dist_esm_index_mjs.6bac1a8c4cc93f15f6b7.js.map +0 -0
  60. {mito_ai-0.1.46.data → mito_ai-0.1.47.data}/data/share/jupyter/labextensions/mito_ai/static/vendors-node_modules_aws-amplify_ui-react_dist_esm_index_mjs.4fcecd65bef9e9847609.js +0 -0
  61. {mito_ai-0.1.46.data → mito_ai-0.1.47.data}/data/share/jupyter/labextensions/mito_ai/static/vendors-node_modules_aws-amplify_ui-react_dist_esm_index_mjs.4fcecd65bef9e9847609.js.map +0 -0
  62. {mito_ai-0.1.46.data → mito_ai-0.1.47.data}/data/share/jupyter/labextensions/mito_ai/static/vendors-node_modules_react-dom_client_js-node_modules_aws-amplify_ui-react_dist_styles_css.b43d4249e4d3dac9ad7b.js +0 -0
  63. {mito_ai-0.1.46.data → mito_ai-0.1.47.data}/data/share/jupyter/labextensions/mito_ai/static/vendors-node_modules_react-dom_client_js-node_modules_aws-amplify_ui-react_dist_styles_css.b43d4249e4d3dac9ad7b.js.map +0 -0
  64. {mito_ai-0.1.46.data → mito_ai-0.1.47.data}/data/share/jupyter/labextensions/mito_ai/static/vendors-node_modules_semver_index_js.3f6754ac5116d47de76b.js +0 -0
  65. {mito_ai-0.1.46.data → mito_ai-0.1.47.data}/data/share/jupyter/labextensions/mito_ai/static/vendors-node_modules_semver_index_js.3f6754ac5116d47de76b.js.map +0 -0
  66. {mito_ai-0.1.46.data → mito_ai-0.1.47.data}/data/share/jupyter/labextensions/mito_ai/static/vendors-node_modules_vscode-diff_dist_index_js.ea55f1f9346638aafbcf.js +0 -0
  67. {mito_ai-0.1.46.data → mito_ai-0.1.47.data}/data/share/jupyter/labextensions/mito_ai/static/vendors-node_modules_vscode-diff_dist_index_js.ea55f1f9346638aafbcf.js.map +0 -0
  68. {mito_ai-0.1.46.dist-info → mito_ai-0.1.47.dist-info}/WHEEL +0 -0
  69. {mito_ai-0.1.46.dist-info → mito_ai-0.1.47.dist-info}/entry_points.txt +0 -0
  70. {mito_ai-0.1.46.dist-info → mito_ai-0.1.47.dist-info}/licenses/LICENSE +0 -0
mito_ai/_version.py CHANGED
@@ -1,4 +1,4 @@
1
1
  # This file is auto-generated by Hatchling. As such, do not:
2
2
  # - modify
3
3
  # - track in version control e.g. be sure to add to .gitignore
4
- __version__ = VERSION = '0.1.46'
4
+ __version__ = VERSION = '0.1.47'
@@ -4,10 +4,11 @@
4
4
  import os
5
5
  import time
6
6
  import logging
7
- from typing import Any, Union, List
7
+ from typing import Any, Union, List, Optional
8
8
  import tempfile
9
- from mito_ai.streamlit_conversion.streamlit_utils import get_app_path
9
+ from mito_ai.path_utils import AbsoluteAppPath, does_app_path_exists, get_absolute_app_path, get_absolute_notebook_dir_path, get_absolute_notebook_path
10
10
  from mito_ai.utils.create import initialize_user
11
+ from mito_ai.utils.error_classes import StreamlitDeploymentError
11
12
  from mito_ai.utils.version_utils import is_pro
12
13
  from mito_ai.utils.websocket_base import BaseWebSocketHandler
13
14
  from mito_ai.app_deploy.app_deploy_utils import add_files_to_zip
@@ -16,11 +17,13 @@ from mito_ai.app_deploy.models import (
16
17
  AppDeployError,
17
18
  DeployAppRequest,
18
19
  ErrorMessage,
19
- MessageType
20
20
  )
21
+ from mito_ai.completions.models import MessageType
21
22
  from mito_ai.logger import get_logger
22
23
  from mito_ai.constants import ACTIVE_STREAMLIT_BASE_URL
24
+ from mito_ai.utils.telemetry_utils import log_streamlit_app_deployment_failure
23
25
  import requests
26
+ import traceback
24
27
 
25
28
 
26
29
  class AppDeployHandler(BaseWebSocketHandler):
@@ -67,41 +70,46 @@ class AppDeployHandler(BaseWebSocketHandler):
67
70
  self.log.debug("App builder message received: %s", message)
68
71
 
69
72
  try:
70
- # Ensure message is a string before parsing
71
- if not isinstance(message, str):
72
- raise ValueError("Message must be a string")
73
-
74
73
  parsed_message = self.parse_message(message)
75
74
  message_type = parsed_message.get('type')
76
75
 
77
76
  if message_type == MessageType.DEPLOY_APP.value:
78
77
  # Handle build app request
79
78
  deploy_app_request = DeployAppRequest(**parsed_message)
80
- await self._handle_deploy_app(deploy_app_request)
79
+ response = await self._handle_deploy_app(deploy_app_request)
80
+ self.reply(response)
81
81
  else:
82
82
  self.log.error(f"Unknown message type: {message_type}")
83
83
  error = AppDeployError(
84
84
  error_type="InvalidRequest",
85
- title=f"Unknown message type: {message_type}"
85
+ message=f"Unknown message type: {message_type}",
86
+ error_code=400
86
87
  )
87
- self.reply(ErrorMessage(**error.__dict__))
88
-
89
- except ValueError as e:
88
+ raise StreamlitDeploymentError(error)
89
+
90
+ except StreamlitDeploymentError as e:
90
91
  self.log.error("Invalid app builder request", exc_info=e)
91
- error = AppDeployError.from_exception(e)
92
- self.reply(ErrorMessage(**error.__dict__))
92
+ log_streamlit_app_deployment_failure('mito_server_key', MessageType.DEPLOY_APP, e.error.__dict__)
93
+ self.reply(
94
+ DeployAppReply(
95
+ parent_id=e.message_id,
96
+ url="",
97
+ error=ErrorMessage(**e.error.__dict__)
98
+ )
99
+ )
93
100
  except Exception as e:
94
101
  self.log.error("Error handling app builder message", exc_info=e)
95
102
  error = AppDeployError.from_exception(
96
103
  e,
97
104
  hint="An error occurred while building the app. Please check the logs for details."
98
105
  )
106
+ log_streamlit_app_deployment_failure('mito_server_key', MessageType.DEPLOY_APP, error.__dict__)
99
107
  self.reply(ErrorMessage(**error.__dict__))
100
108
 
101
109
  latency_ms = round((time.time() - start) * 1000)
102
110
  self.log.info(f"App builder handler processed in {latency_ms} ms.")
103
111
 
104
- async def _handle_deploy_app(self, message: DeployAppRequest) -> None:
112
+ async def _handle_deploy_app(self, message: DeployAppRequest) -> DeployAppReply:
105
113
  """Handle a build app request.
106
114
 
107
115
  Args:
@@ -114,19 +122,22 @@ class AppDeployHandler(BaseWebSocketHandler):
114
122
 
115
123
  if not message_id:
116
124
  self.log.error("Missing message_id in request")
117
- return
118
-
125
+ error = AppDeployError(
126
+ error_type="BadRequest",
127
+ message="Missing message_id in request",
128
+ error_code=400,
129
+ message_id=message_id
130
+ )
131
+ raise StreamlitDeploymentError(error)
132
+
119
133
  if not notebook_path:
120
134
  error = AppDeployError(
121
135
  error_type="InvalidRequest",
122
- title="Missing 'notebook_path' parameter"
136
+ message="Missing 'notebook_path' parameter",
137
+ error_code=400,
138
+ message_id=message_id
123
139
  )
124
- self.reply(DeployAppReply(
125
- parent_id=message_id,
126
- url="",
127
- error=error
128
- ))
129
- return
140
+ raise StreamlitDeploymentError(error)
130
141
 
131
142
  # Validate JWT token if provided
132
143
  token_preview = jwt_token[:20] if jwt_token else "No token provided"
@@ -136,55 +147,42 @@ class AppDeployHandler(BaseWebSocketHandler):
136
147
  self.log.error("JWT token validation failed")
137
148
  error = AppDeployError(
138
149
  error_type="Unauthorized",
139
- title="Invalid authentication token",
140
- hint="Please sign in again to deploy your app."
150
+ message="Invalid authentication token",
151
+ hint="Please sign in again to deploy your app.",
152
+ error_code=401,
153
+ message_id=message_id
141
154
  )
142
- self.reply(DeployAppReply(
143
- parent_id=message_id,
144
- url="",
145
- error=error
146
- ))
147
- return
155
+ raise StreamlitDeploymentError(error)
148
156
  else:
149
157
  self.log.info("JWT token validation successful")
150
-
151
- try:
152
- notebook_path = str(notebook_path) if notebook_path else ""
153
158
 
154
- app_directory = os.path.dirname(notebook_path)
155
-
156
- # Check if the app.py file exists
157
- app_path = get_app_path(app_directory)
158
- if app_path is None:
159
- error = AppDeployError(
160
- error_type="AppNotFound",
161
- title="App not found",
162
- hint="Please make sure the app.py file exists in the same directory as the notebook."
163
- )
164
- self.reply(DeployAppReply(
165
- parent_id=message_id,
166
- url="",
167
- error=error
168
- ))
169
-
170
- # Finally, deploy the app
171
- deploy_url = await self._deploy_app(app_directory, files_to_upload, jwt_token)
159
+ notebook_path = str(notebook_path) if notebook_path else ""
160
+ absolute_notebook_path = get_absolute_notebook_path(notebook_path)
161
+ absolute_app_directory = get_absolute_notebook_dir_path(absolute_notebook_path)
162
+ app_path = get_absolute_app_path(absolute_app_directory)
163
+
164
+ # Check if the app.py file exists
165
+ app_path_exists = does_app_path_exists(app_path)
166
+ if not app_path_exists:
167
+ error = AppDeployError(
168
+ error_type="AppNotFound",
169
+ message="App not found",
170
+ hint="Please make sure the app.py file exists in the same directory as the notebook.",
171
+ error_code=400,
172
+ message_id=message_id
173
+ )
174
+ raise StreamlitDeploymentError(error)
175
+
176
+ # Finally, deploy the app
177
+ deploy_url = await self._deploy_app(app_path, files_to_upload, message_id, jwt_token)
178
+
179
+ # Send the response
180
+ return DeployAppReply(
181
+ parent_id=message_id,
182
+ url=deploy_url if deploy_url else ""
183
+ )
184
+
172
185
 
173
- # Send the response
174
- self.reply(DeployAppReply(
175
- parent_id=message_id,
176
- url=deploy_url
177
- ))
178
-
179
- except Exception as e:
180
- self.log.error(f"Error building app: {e}", exc_info=e)
181
- error = AppDeployError.from_exception(e)
182
- self.reply(DeployAppReply(
183
- parent_id=message_id,
184
- url="",
185
- error=error
186
- ))
187
-
188
186
  def _validate_jwt_token(self, token: str) -> bool:
189
187
  """Basic JWT token validation logic.
190
188
 
@@ -219,7 +217,7 @@ class AppDeployHandler(BaseWebSocketHandler):
219
217
  return False
220
218
 
221
219
 
222
- async def _deploy_app(self, app_path: str, files_to_upload:List[str], jwt_token: str = '') -> str:
220
+ async def _deploy_app(self, app_path: AbsoluteAppPath, files_to_upload:List[str], message_id: str, jwt_token: str = '') -> Optional[str]:
223
221
  """Deploy the app using pre-signed URLs.
224
222
 
225
223
  Args:
@@ -255,7 +253,7 @@ class AppDeployHandler(BaseWebSocketHandler):
255
253
  expected_app_url = url_data['expected_app_url']
256
254
 
257
255
  self.log.info(f"Received pre-signed URL. App will be available at: {expected_app_url}")
258
-
256
+
259
257
  # Step 2: Create a zip file of the app.
260
258
  temp_zip_path = None
261
259
  try:
@@ -269,7 +267,14 @@ class AppDeployHandler(BaseWebSocketHandler):
269
267
  upload_response = await self._upload_app_to_s3(temp_zip_path, presigned_url)
270
268
  except Exception as e:
271
269
  self.log.error(f"Error zipping app: {e}")
272
- raise
270
+ error = AppDeployError(
271
+ error_type="ZippingError",
272
+ message=f"Error zipping app: {e}",
273
+ traceback=traceback.format_exc(),
274
+ error_code=500,
275
+ message_id=message_id
276
+ )
277
+ raise StreamlitDeploymentError(error)
273
278
  finally:
274
279
  # Clean up
275
280
  if temp_zip_path is not None:
@@ -278,20 +283,35 @@ class AppDeployHandler(BaseWebSocketHandler):
278
283
  self.log.info(f"Upload successful! Status code: {upload_response.status_code}")
279
284
 
280
285
  self.log.info(f"Deployment initiated. App will be available at: {expected_app_url}")
281
- return expected_app_url # type: ignore
286
+ return str(expected_app_url)
282
287
 
283
288
  except requests.exceptions.RequestException as e:
284
289
  self.log.error(f"Error during API request: {e}")
285
290
  if hasattr(e, 'response') and e.response is not None:
286
291
  error_detail = e.response.json()
292
+ error_message = error_detail.get('error', "")
287
293
  self.log.error(f"Server error details: {error_detail}")
288
- if 'error' in error_detail:
289
- raise Exception(error_detail['error'])
290
- raise
294
+ else:
295
+ error_message = str(e)
296
+
297
+ error = AppDeployError(
298
+ error_type="APIException",
299
+ message=str(error_message),
300
+ traceback=traceback.format_exc(),
301
+ error_code=500,
302
+ message_id=message_id
303
+ )
304
+ raise StreamlitDeploymentError(error)
291
305
  except Exception as e:
292
306
  self.log.error(f"Error during deployment: {str(e)}")
293
- raise
294
- raise RuntimeError("Unexpected error in _deploy_app")
307
+ error = AppDeployError(
308
+ error_type="DeploymentException",
309
+ message=str(e),
310
+ traceback=traceback.format_exc(),
311
+ error_code=500,
312
+ message_id=message_id
313
+ )
314
+ raise StreamlitDeploymentError(error)
295
315
 
296
316
  async def _upload_app_to_s3(self, app_path: str, presigned_url: str) -> requests.Response:
297
317
  """Upload the app to S3 using the presigned URL."""
@@ -2,15 +2,10 @@
2
2
  # Distributed under the terms of the GNU Affero General Public License v3.0 License.
3
3
 
4
4
  from dataclasses import dataclass
5
- from enum import Enum
5
+ import traceback
6
6
  from typing import Literal, Optional, List
7
7
 
8
8
 
9
- class MessageType(str, Enum):
10
- """Types of app deploy messages."""
11
- DEPLOY_APP = "deploy-app"
12
-
13
-
14
9
  @dataclass(frozen=True)
15
10
  class AppDeployError:
16
11
  """Error information for app deploy operations."""
@@ -19,7 +14,13 @@ class AppDeployError:
19
14
  error_type: str
20
15
 
21
16
  # Error title.
22
- title: str
17
+ message: str
18
+
19
+ #ID of parent to resolve response
20
+ message_id: Optional[str] = "InvalidMessageID"
21
+
22
+ # Error code
23
+ error_code: Optional[int] = 500
23
24
 
24
25
  # Error traceback.
25
26
  traceback: Optional[str] = None
@@ -28,21 +29,24 @@ class AppDeployError:
28
29
  hint: Optional[str] = None
29
30
 
30
31
  @classmethod
31
- def from_exception(cls, e: Exception, hint: Optional[str] = None) -> "AppDeployError":
32
+ def from_exception(cls, e: Exception, hint: Optional[str] = None, error_code: Optional[int] = 500) -> "AppDeployError":
32
33
  """Create an error from an exception.
33
34
 
34
35
  Args:
35
36
  e: The exception.
36
37
  hint: Optional hint to fix the error.
38
+ error_code: Optional error code which defaults to 500
37
39
 
38
40
  Returns:
39
41
  The app builder error.
40
42
  """
43
+ tb_str = "".join(traceback.format_exception(type(e), e, e.__traceback__))
41
44
  return cls(
42
45
  error_type=type(e).__name__,
43
- title=str(e),
44
- traceback=getattr(e, "__traceback__", None) and str(e.__traceback__),
46
+ message=str(e),
47
+ traceback=tb_str,
45
48
  hint=hint,
49
+ error_code=error_code
46
50
  )
47
51
 
48
52
 
@@ -59,7 +63,7 @@ class DeployAppRequest:
59
63
  """Request to deploy an app."""
60
64
 
61
65
  # Request type.
62
- type: Literal["deploy-app"]
66
+ type: Literal["deploy_app"]
63
67
 
64
68
  # Message ID.
65
69
  message_id: str
@@ -88,4 +92,4 @@ class DeployAppReply:
88
92
  error: Optional[AppDeployError] = None
89
93
 
90
94
  # Type of reply.
91
- type: Literal["deploy-app"] = "deploy-app"
95
+ type: Literal["deploy_app"] = "deploy_app"
@@ -29,12 +29,13 @@ class CellUpdate(BaseModel):
29
29
  # for now and rely on the AI to respond with the correct types, following the format
30
30
  # that we show it in the system prompt.
31
31
  class AgentResponse(BaseModel):
32
- type: Literal['cell_update', 'get_cell_output', 'run_all_cells', 'finished_task']
32
+ type: Literal['cell_update', 'get_cell_output', 'run_all_cells', 'finished_task', 'create_streamlit_app', 'edit_streamlit_app']
33
33
  message: str
34
34
  cell_update: Optional[CellUpdate]
35
35
  get_cell_output_cell_id: Optional[str]
36
36
  next_steps: Optional[List[str]]
37
37
  analysis_assumptions: Optional[List[str]]
38
+ edit_streamlit_app_prompt: Optional[str]
38
39
 
39
40
 
40
41
  @dataclass(frozen=True)
@@ -65,6 +66,7 @@ class MessageType(Enum):
65
66
  UPDATE_MODEL_CONFIG = "update_model_config"
66
67
  STREAMLIT_CONVERSION = "streamlit_conversion"
67
68
  STOP_AGENT = "stop_agent"
69
+ DEPLOY_APP = "deploy_app"
68
70
 
69
71
 
70
72
  @dataclass(frozen=True)
@@ -101,6 +103,7 @@ class AgentExecutionMetadata():
101
103
  files: Optional[List[str]] = None
102
104
  index: Optional[int] = None
103
105
  additionalContext: Optional[List[Dict[str, str]]] = None
106
+ streamlitAppIsOpen: Optional[bool] = None
104
107
 
105
108
  @dataclass(frozen=True)
106
109
  class AgentSmartDebugMetadata():
@@ -5,12 +5,14 @@ from mito_ai.completions.models import AgentExecutionMetadata
5
5
  from mito_ai.completions.prompt_builders.prompt_constants import (
6
6
  FILES_SECTION_HEADING,
7
7
  JUPYTER_NOTEBOOK_SECTION_HEADING,
8
+ STREAMLIT_APP_STATUS_SECTION_HEADING,
8
9
  VARIABLES_SECTION_HEADING,
9
10
  cell_update_output_str
10
11
  )
11
12
  from mito_ai.completions.prompt_builders.utils import (
12
13
  get_rules_str,
13
14
  get_selected_context_str,
15
+ get_streamlit_app_status_str,
14
16
  )
15
17
 
16
18
 
@@ -20,7 +22,7 @@ def create_agent_execution_prompt(md: AgentExecutionMetadata) -> str:
20
22
  ai_optimized_cells_str = '\n'.join([f"{cell}" for cell in md.aiOptimizedCells or []])
21
23
  rules_str = get_rules_str(md.additionalContext)
22
24
  selected_context_str = get_selected_context_str(md.additionalContext)
23
-
25
+ streamlit_status_str = get_streamlit_app_status_str(md.streamlitAppIsOpen)
24
26
  context_str = f"""Remember to choose the correct tool to respond with.
25
27
 
26
28
  {rules_str}
@@ -35,6 +37,9 @@ def create_agent_execution_prompt(md: AgentExecutionMetadata) -> str:
35
37
  {FILES_SECTION_HEADING}
36
38
  {files_str}
37
39
 
40
+ {STREAMLIT_APP_STATUS_SECTION_HEADING}
41
+ {streamlit_status_str}
42
+
38
43
  {selected_context_str}
39
44
 
40
45
  {cell_update_output_str(md.base64EncodedActiveCellOutput is not None)}"""
@@ -225,6 +225,57 @@ Important information:
225
225
  5. Do not use this tool repeatedly if it continues to produce errors - instead, focus on fixing the specific error that occurred.
226
226
  ====
227
227
 
228
+ TOOL: CREATE_STREAMLIT_APP
229
+
230
+ When you want to create a new Streamlit app from the current notebook, respond with this format:
231
+
232
+ {{
233
+ type: 'create_streamlit_app',
234
+ message: str
235
+ }}
236
+
237
+ Important information:
238
+ 1. The message is a short summary of why you're creating the Streamlit app.
239
+ 2. Only use this tool when the user explicitly asks to create or preview a Streamlit app AND no Streamlit app is currently open.
240
+ 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.
241
+ 4. Using this tool will automatically open the app so the user can see a preview of the app.
242
+ 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.
243
+
244
+ <Example>
245
+
246
+ Your task: Show me my notebook as an app.
247
+
248
+ Output:
249
+ {{
250
+ type: 'create_streamlit_app',
251
+ message: "I'll convert your notebook into an app."
252
+ }}
253
+
254
+ 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.
255
+
256
+ <Example>
257
+
258
+ ====
259
+
260
+ TOOL: EDIT_STREAMLIT_APP
261
+
262
+ When you want to edit an existing Streamlit app, respond with this format:
263
+
264
+ {{
265
+ type: 'edit_streamlit_app',
266
+ message: str,
267
+ edit_streamlit_app_prompt: str
268
+ }}
269
+
270
+ Important information:
271
+ 1. The message is a short summary of why you're editing the Streamlit app.
272
+ 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").
273
+ 3. Only use this tool when the user asks to edit, update, or modify a Streamlit app.
274
+ 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.
275
+ 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.
276
+
277
+ ====
278
+
228
279
  TOOL: FINISHED_TASK
229
280
 
230
281
  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:
238
289
  Important information:
239
290
  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.""
240
291
  2. The message should include citations for any insights that you shared with the user.
241
- 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"].
242
- 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"].
292
+ 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"].
293
+ 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"].
243
294
  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.
244
295
  6. If the user's task doesn't involve creating or modifying a code cell, you should respond with a FINISHED_TASK response.
245
296
  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
391
442
 
392
443
  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.
393
444
 
394
- 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.
445
+ 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.
446
+
447
+ 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.
395
448
 
396
449
  Whenever you get a message back from the user, you should:
397
450
  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
410
463
  - 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.
411
464
  - Wait for the user to send you back the updated variables and notebook state.
412
465
  {'' 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.'}
413
- - If after reviewing the updates you decide that you've completed the task, send a FINISHED_TASK tool message."""
466
+ - If after reviewing the updates you decide that you've completed the task, send a FINISHED_TASK tool message.
467
+
468
+ ====
469
+
470
+ OTHER USEFUL INFORMATION:
471
+ 1. When importing matplotlib, write the code `%matplotlib inline` to make sure the graphs render in Jupyter
472
+ """
@@ -21,6 +21,10 @@ There are three possible types of responses you might give:
21
21
  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.
22
22
  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.
23
23
 
24
+ Other useful information:
25
+ 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.
26
+ 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.
27
+
24
28
  ====
25
29
  {CITATION_RULES}
26
30
 
@@ -19,6 +19,7 @@ ACTIVE_CELL_ID_SECTION_HEADING = "The ID of the active code cell:"
19
19
  ACTIVE_CELL_OUTPUT_SECTION_HEADING = "Output of the active code cell:"
20
20
  GET_CELL_OUTPUT_TOOL_RESPONSE_SECTION_HEADING = "Output of the code cell you just applied the CELL_UPDATE to:"
21
21
  JUPYTER_NOTEBOOK_SECTION_HEADING = "Jupyter Notebook:"
22
+ STREAMLIT_APP_STATUS_SECTION_HEADING = "Streamlit App Status:"
22
23
 
23
24
  # Placeholder text used when trimming content from messages
24
25
  CONTENT_REMOVED_PLACEHOLDER = "Content removed to save space"
@@ -69,3 +69,17 @@ def get_selected_context_str(additional_context: Optional[List[Dict[str, str]]])
69
69
 
70
70
  # STEP 3: Combine into a single string
71
71
  return "\n\n".join(context_parts)
72
+
73
+
74
+ def get_streamlit_app_status_str(streamlit_app_is_open: Optional[bool]) -> str:
75
+ """
76
+ Get the streamlit app status string.
77
+ """
78
+ if streamlit_app_is_open is None:
79
+ return ""
80
+
81
+ if streamlit_app_is_open:
82
+ return "A Streamlit app is currently open and running."
83
+ else:
84
+ return "No Streamlit app is currently open."
85
+
mito_ai/path_utils.py ADDED
@@ -0,0 +1,56 @@
1
+ # Copyright (c) Saga Inc.
2
+ # Distributed under the terms of the GNU Affero General Public License v3.0 License.
3
+
4
+ from typing import NewType
5
+ import os
6
+ from mito_ai.utils.error_classes import StreamlitPreviewError
7
+
8
+ # Type definitions for better type safety
9
+ AbsoluteNotebookPath = NewType('AbsoluteNotebookPath', str)
10
+ AbsoluteNotebookDirPath = NewType('AbsoluteNotebookDirPath', str)
11
+ AbsoluteAppPath = NewType('AbsoluteAppPath', str)
12
+
13
+ def get_absolute_notebook_path(notebook_path: str) -> AbsoluteNotebookPath:
14
+ """
15
+ Convert any notebook path to an absolute path.
16
+
17
+ Args:
18
+ notebook_path: Path to the notebook (can be relative or absolute)
19
+
20
+ Returns:
21
+ AbsoluteNotebookPath: The absolute path to the notebook
22
+
23
+ Raises:
24
+ ValueError: If the path is invalid or empty
25
+ """
26
+ if not notebook_path or not notebook_path.strip():
27
+ raise StreamlitPreviewError("Notebook path cannot be empty", 400)
28
+
29
+ absolute_path = os.path.abspath(notebook_path)
30
+ return AbsoluteNotebookPath(absolute_path)
31
+
32
+
33
+ def get_absolute_notebook_dir_path(notebook_path: AbsoluteNotebookPath) -> AbsoluteNotebookDirPath:
34
+ """
35
+ Get the directory containing the notebook.
36
+
37
+ Args:
38
+ notebook_path: Absolute path to the notebook
39
+
40
+ Returns:
41
+ AbsoluteNotebookDirPath: The directory containing the notebook
42
+ """
43
+ return AbsoluteNotebookDirPath(os.path.dirname(notebook_path))
44
+
45
+ def get_absolute_app_path(app_directory: AbsoluteNotebookDirPath) -> AbsoluteAppPath:
46
+ """
47
+ Check if the app.py file exists in the given directory.
48
+ """
49
+ return AbsoluteAppPath(os.path.join(app_directory, "app.py"))
50
+
51
+ def does_app_path_exists(app_path: AbsoluteAppPath) -> bool:
52
+ """
53
+ Check if the app.py file exists in the given directory.
54
+ """
55
+ return os.path.exists(app_path)
56
+