mito-ai 0.1.42__tar.gz → 0.1.44__tar.gz

{mito_ai-0.1.42 → mito_ai-0.1.44}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: mito_ai
-Version: 0.1.42
+Version: 0.1.44
 Summary: AI chat for JupyterLab
 Project-URL: Homepage, https://trymito.io
 Project-URL: Bug Tracker, https://github.com/mito-ds/monorepo/issues

{mito_ai-0.1.42 → mito_ai-0.1.44}/mito_ai/__init__.py

@@ -5,7 +5,7 @@ from typing import List, Dict
 from jupyter_server.utils import url_path_join
 from mito_ai.completions.handlers import CompletionHandler
 from mito_ai.completions.providers import OpenAIProvider
-from mito_ai.app_builder.handlers import AppBuilderHandler
+from mito_ai.app_deploy.handlers import AppDeployHandler
 from mito_ai.streamlit_preview.handlers import StreamlitPreviewHandler
 from mito_ai.log.urls import get_log_urls
 from mito_ai.version_check import VersionCheckHandler
@@ -71,8 +71,8 @@ def _load_jupyter_server_extension(server_app) -> None: # type: ignore
             {"llm": open_ai_provider},
         ),
         (
-            url_path_join(base_url, "mito-ai", "app-builder"),
-            AppBuilderHandler,
+            url_path_join(base_url, "mito-ai", "app-deploy"),
+            AppDeployHandler,
             {}
         ),
         (

{mito_ai-0.1.42 → mito_ai-0.1.44}/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.42'
+__version__ = VERSION = '0.1.44'

{mito_ai-0.1.42/mito_ai/app_builder → mito_ai-0.1.44/mito_ai/app_deploy}/__init__.py

@@ -3,4 +3,4 @@
 
 """App builder module for Mito AI."""
 
-from .handlers import AppBuilderHandler 
+from .handlers import AppDeployHandler 

{mito_ai-0.1.42/mito_ai/app_builder → mito_ai-0.1.44/mito_ai/app_deploy}/handlers.py

@@ -7,13 +7,14 @@ import logging
 from typing import Any, Union, Optional
 import zipfile
 import tempfile
+from mito_ai.streamlit_conversion.streamlit_utils import get_app_path
 from mito_ai.utils.create import initialize_user
 from mito_ai.utils.version_utils import is_pro
 from mito_ai.utils.websocket_base import BaseWebSocketHandler
-from mito_ai.app_builder.models import (
-    BuildAppReply,
-    AppBuilderError,
-    BuildAppRequest,
+from mito_ai.app_deploy.models import (
+    DeployAppReply,
+    AppDeployError,
+    DeployAppRequest,
     ErrorMessage,
     MessageType
 )
@@ -23,8 +24,8 @@ from mito_ai.constants import ACTIVE_STREAMLIT_BASE_URL
 import requests
 
 
-class AppBuilderHandler(BaseWebSocketHandler):
-    """Handler for app building requests."""
+class AppDeployHandler(BaseWebSocketHandler):
+    """Handler for app deploy requests."""
     
     def initialize(self) -> None:
         """Initialize the WebSocket handler."""
@@ -57,6 +58,7 @@ class AppBuilderHandler(BaseWebSocketHandler):
         Args:
             message: The message received on the WebSocket.
         """
+        
         start = time.time()
         
         # Convert bytes to string if needed
@@ -73,13 +75,13 @@ class AppBuilderHandler(BaseWebSocketHandler):
             parsed_message = self.parse_message(message)
             message_type = parsed_message.get('type')
             
-            if message_type == MessageType.BUILD_APP.value:
+            if message_type == MessageType.DEPLOY_APP.value:
                 # Handle build app request
-                build_app_request = BuildAppRequest(**parsed_message)
-                await self._handle_build_app(build_app_request)
+                deploy_app_request = DeployAppRequest(**parsed_message)
+                await self._handle_deploy_app(deploy_app_request)
             else:
                 self.log.error(f"Unknown message type: {message_type}")
-                error = AppBuilderError(
+                error = AppDeployError(
                     error_type="InvalidRequest",
                     title=f"Unknown message type: {message_type}"
                 )
@@ -87,11 +89,11 @@ class AppBuilderHandler(BaseWebSocketHandler):
                 
         except ValueError as e:
             self.log.error("Invalid app builder request", exc_info=e)
-            error = AppBuilderError.from_exception(e)
+            error = AppDeployError.from_exception(e)
             self.reply(ErrorMessage(**error.__dict__))
         except Exception as e:
             self.log.error("Error handling app builder message", exc_info=e)
-            error = AppBuilderError.from_exception(
+            error = AppDeployError.from_exception(
                 e, 
                 hint="An error occurred while building the app. Please check the logs for details."
             )
@@ -100,7 +102,7 @@ class AppBuilderHandler(BaseWebSocketHandler):
         latency_ms = round((time.time() - start) * 1000)
         self.log.info(f"App builder handler processed in {latency_ms} ms.")
     
-    async def _handle_build_app(self, message: BuildAppRequest) -> None:
+    async def _handle_deploy_app(self, message: DeployAppRequest) -> None:
         """Handle a build app request.
         
         Args:
@@ -109,17 +111,17 @@ class AppBuilderHandler(BaseWebSocketHandler):
         message_id = message.message_id
         notebook_path = message.notebook_path
         jwt_token = message.jwt_token
-
+        
         if not message_id:
             self.log.error("Missing message_id in request")
             return
         
         if not notebook_path:
-            error = AppBuilderError(
+            error = AppDeployError(
                 error_type="InvalidRequest",
                 title="Missing 'notebook_path' parameter"
             )
-            self.reply(BuildAppReply(
+            self.reply(DeployAppReply(
                 parent_id=message_id,
                 url="",
                 error=error
@@ -132,12 +134,12 @@ class AppBuilderHandler(BaseWebSocketHandler):
         is_valid = self._validate_jwt_token(jwt_token) if jwt_token else False
         if not is_valid or not jwt_token:
             self.log.error("JWT token validation failed")
-            error = AppBuilderError(
+            error = AppDeployError(
                 error_type="Unauthorized",
                 title="Invalid authentication token",
                 hint="Please sign in again to deploy your app."
             )
-            self.reply(BuildAppReply(
+            self.reply(DeployAppReply(
                 parent_id=message_id,
                 url="",
                 error=error
@@ -150,25 +152,34 @@ class AppBuilderHandler(BaseWebSocketHandler):
             notebook_path = str(notebook_path) if notebook_path else ""
 
             app_directory = os.path.dirname(notebook_path)
-            app_path = os.path.join(app_directory, "app.py")
-
-            if not os.path.exists(app_path):
-                success_flag, app_path_result, result_message = await streamlit_handler(notebook_path)
-                if not success_flag or app_path_result is None:
-                    raise Exception(result_message)
-
+            
+            # 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, jwt_token)
 
             # Send the response
-            self.reply(BuildAppReply(
+            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 = AppBuilderError.from_exception(e)
-            self.reply(BuildAppReply(
+            error = AppDeployError.from_exception(e)
+            self.reply(DeployAppReply(
                 parent_id=message_id,
                 url="",
                 error=error

{mito_ai-0.1.42/mito_ai/app_builder → mito_ai-0.1.44/mito_ai/app_deploy}/models.py

@@ -7,13 +7,13 @@ from typing import Literal, Optional
 
 
 class MessageType(str, Enum):
-    """Types of app builder messages."""
-    BUILD_APP = "build-app"
+    """Types of app deploy messages."""
+    DEPLOY_APP = "deploy-app"
 
 
 @dataclass(frozen=True)
-class AppBuilderError:
-    """Error information for app builder operations."""
+class AppDeployError:
+    """Error information for app deploy operations."""
     
     # Error type.
     error_type: str
@@ -28,7 +28,7 @@ class AppBuilderError:
     hint: Optional[str] = None
     
     @classmethod
-    def from_exception(cls, e: Exception, hint: Optional[str] = None) -> "AppBuilderError":
+    def from_exception(cls, e: Exception, hint: Optional[str] = None) -> "AppDeployError":
         """Create an error from an exception.
 
         Args:
@@ -47,7 +47,7 @@ class AppBuilderError:
 
 
 @dataclass(frozen=True)
-class ErrorMessage(AppBuilderError):
+class ErrorMessage(AppDeployError):
     """Error message."""
     
     # Message type.
@@ -55,11 +55,11 @@ class ErrorMessage(AppBuilderError):
 
 
 @dataclass(frozen=True)
-class BuildAppRequest:
-    """Request to build an app."""
+class DeployAppRequest:
+    """Request to deploy an app."""
     
     # Request type.
-    type: Literal["build-app"]
+    type: Literal["deploy-app"]
     
     # Message ID.
     message_id: str
@@ -72,8 +72,8 @@ class BuildAppRequest:
 
 
 @dataclass(frozen=True)
-class BuildAppReply:
-    """Reply to a build app request."""
+class DeployAppReply:
+    """Reply to a deplpy app request."""
     
     # ID of the request message this is replying to.
     parent_id: str
@@ -82,7 +82,7 @@ class BuildAppReply:
     url: str
     
     # Optional error information.
-    error: Optional[AppBuilderError] = None
+    error: Optional[AppDeployError] = None
     
     # Type of reply.
-    type: Literal["build-app"] = "build-app" 
+    type: Literal["deploy-app"] = "deploy-app" 

{mito_ai-0.1.42 → mito_ai-0.1.44}/mito_ai/app_manager/handlers.py

@@ -12,11 +12,14 @@ from mito_ai.app_manager.models import (
     AppManagerError,
     ManageAppRequest,
     ManageAppReply,
+    CheckAppStatusRequest,
+    CheckAppStatusReply,
     ErrorMessage,
     MessageType
 )
 from mito_ai.constants import ACTIVE_STREAMLIT_BASE_URL
 from mito_ai.logger import get_logger
+from mito_ai.app_manager.utils import convert_utc_to_local_time
 import requests
 
 
@@ -56,6 +59,10 @@ class AppManagerHandler(BaseWebSocketHandler):
                 # Handle manage app request
                 manage_app_request = ManageAppRequest(**parsed_message)
                 await self._handle_manage_app(manage_app_request)
+            elif message_type == MessageType.CHECK_APP_STATUS.value:
+                # Handle check app status request
+                check_status_request = CheckAppStatusRequest(**parsed_message)
+                await self._handle_check_app_status(check_status_request)
             else:
                 self.log.error(f"Unknown message type: {message_type}")
                 error_response = ErrorMessage(
@@ -102,6 +109,10 @@ class AppManagerHandler(BaseWebSocketHandler):
 
             apps_data = manage_apps_response.json()
 
+            for app in apps_data:
+                if 'last_deployed_at' in app:
+                    app['last_deployed_at'] = convert_utc_to_local_time(app['last_deployed_at'])
+
             # Create successful response
             reply = ManageAppReply(
                 apps=apps_data,
@@ -126,4 +137,31 @@ class AppManagerHandler(BaseWebSocketHandler):
                 error=error,
                 message_id=request.message_id
             )
+            self.reply(error_reply)
+
+    async def _handle_check_app_status(self, request: CheckAppStatusRequest) -> None:
+        """Handle a check app status request."""
+        self.log.info("In check app status")
+        try:
+            # Make a HEAD request to check if the app URL is accessible
+            response = requests.head(request.app_url, timeout=10, verify=False)
+            self.log.debug(f"Is app accessible: {response.status_code}")
+            is_accessible = response.status_code==200
+
+            # Create successful response
+            reply = CheckAppStatusReply(
+                is_accessible=is_accessible
+            )
+
+            self.reply(reply)
+
+        except Exception as e:
+            self.log.error(f"Error checking app status: {e}", exc_info=e)
+            error = AppManagerError.from_exception(e)
+
+            # Return error response
+            error_reply = CheckAppStatusReply(
+                is_accessible=False,
+                error=error
+            )
             self.reply(error_reply)

{mito_ai-0.1.42 → mito_ai-0.1.44}/mito_ai/app_manager/models.py

@@ -1,5 +1,4 @@
 # Copyright (c) Saga Inc.
-
 # Distributed under the terms of the GNU Affero General Public License v3.0 License.
 
 from dataclasses import dataclass, field
@@ -9,6 +8,7 @@ from typing import List, Optional
 class MessageType(str, Enum):
     """Types of app manager messages."""
     MANAGE_APP = "manage-app"
+    CHECK_APP_STATUS = "check-app-status"
 
 
 @dataclass(frozen=True)
@@ -16,7 +16,7 @@ class ManageAppRequest:
     """Request to manage apps."""
     type: str = "manage-app"
     jwt_token: Optional[str] = None
-    message_id: Optional[str] = None  
+    message_id: Optional[str] = None
 
 @dataclass(frozen=True)
 class App:
@@ -24,7 +24,7 @@ class App:
     app_name: str
     url: str
     status: str
-    created_at: str
+    last_deployed_at: str
 
 @dataclass(frozen=True)
 class AppManagerError:
@@ -49,6 +49,19 @@ class ManageAppReply:
     error: Optional[AppManagerError] = None
     message_id: Optional[str] = None  
 
+@dataclass(frozen=True)
+class CheckAppStatusRequest:
+    """Request to check app status."""
+    app_url: str
+    type: str = "check-app-status"
+
+@dataclass(frozen=True)
+class CheckAppStatusReply:
+    """Reply to a check app status request."""
+    is_accessible: bool
+    type: str = "check-app-status"
+    error: Optional[AppManagerError] = None
+
 @dataclass(frozen=True)
 class ErrorMessage:
     """Error message."""

mito_ai-0.1.44/mito_ai/app_manager/utils.py

@@ -0,0 +1,24 @@
+# Copyright (c) Saga Inc.
+# Distributed under the terms of the GNU Affero General Public License v3.0 License.
+
+from datetime import datetime, timezone
+
+def convert_utc_to_local_time(time_str: str) -> str:
+    """Convert UTC time to a user's local time"""
+    try:
+        # Remove the 'Z' suffix and parse the UTC datetime
+        utc_time_str = time_str.rstrip('Z')
+        utc_time = datetime.fromisoformat(utc_time_str)
+
+        # Set timezone to UTC
+        utc_time = utc_time.replace(tzinfo=timezone.utc)
+
+        # Convert to local timezone (system timezone)
+        local_time = utc_time.astimezone()
+
+        # Format as 'MMM DD HH:MM'
+        return local_time.strftime('%m-%d-%Y %H:%M')
+
+    except (ValueError, AttributeError) as e:
+        # Return original string if parsing fails
+        return time_str

{mito_ai-0.1.42 → mito_ai-0.1.44}/mito_ai/completions/completion_handlers/agent_execution_handler.py

@@ -38,7 +38,7 @@ class AgentExecutionHandler(CompletionHandler[AgentExecutionMetadata]):
         display_prompt = metadata.input
         
         # Add the prompt to the message history
-        new_ai_optimized_message = create_ai_optimized_message(prompt, metadata.base64EncodedActiveCellOutput, metadata.base64EncodedUploadedImage)
+        new_ai_optimized_message = create_ai_optimized_message(prompt, metadata.base64EncodedActiveCellOutput, metadata.additionalContext)
         new_display_optimized_message: ChatCompletionMessageParam = {"role": "user", "content": display_prompt}
         
         await message_history.append_message(new_ai_optimized_message, new_display_optimized_message, model, provider, metadata.threadId)

{mito_ai-0.1.42 → mito_ai-0.1.44}/mito_ai/completions/completion_handlers/chat_completion_handler.py

@@ -47,7 +47,7 @@ class ChatCompletionHandler(CompletionHandler[ChatMessageMetadata]):
         display_prompt = f"```python{metadata.activeCellCode or ''}```{metadata.input}"
         
         # Add the prompt to the message history
-        new_ai_optimized_message = create_ai_optimized_message(prompt, metadata.base64EncodedActiveCellOutput, metadata.base64EncodedUploadedImage)
+        new_ai_optimized_message = create_ai_optimized_message(prompt, metadata.base64EncodedActiveCellOutput, metadata.additionalContext)
         new_display_optimized_message: ChatCompletionMessageParam = {"role": "user", "content": display_prompt}
         await message_history.append_message(new_ai_optimized_message, new_display_optimized_message, model, provider, metadata.threadId)
         
@@ -110,7 +110,7 @@ class ChatCompletionHandler(CompletionHandler[ChatMessageMetadata]):
         display_prompt = f"```python{metadata.activeCellCode or ''}```{metadata.input}"
         
         # Add the prompt to the message history
-        new_ai_optimized_message = create_ai_optimized_message(prompt, metadata.base64EncodedActiveCellOutput, metadata.base64EncodedUploadedImage)
+        new_ai_optimized_message = create_ai_optimized_message(prompt, metadata.base64EncodedActiveCellOutput, metadata.additionalContext)
         new_display_optimized_message: ChatCompletionMessageParam = {"role": "user", "content": display_prompt}
         await message_history.append_message(new_ai_optimized_message, new_display_optimized_message, model, provider, metadata.threadId)
         

{mito_ai-0.1.42 → mito_ai-0.1.44}/mito_ai/completions/completion_handlers/utils.py

@@ -1,6 +1,7 @@
 # Copyright (c) Saga Inc.
 # Distributed under the terms of the GNU Affero General Public License v3.0 License.
 
+import base64
 from typing import Optional, Union, List, Dict, Any, cast
 from mito_ai.completions.message_history import GlobalMessageHistory
 from mito_ai.completions.models import ThreadID
@@ -75,16 +76,37 @@ async def append_agent_system_message(
     )
 
 
+def extract_and_encode_images_from_additional_context(
+    additional_context: Optional[List[Dict[str, str]]],
+) -> List[str]:
+    encoded_images = []
+
+    for context in additional_context or []:
+        if context["type"].startswith("image/"):
+            try:
+                with open(context["value"], "rb") as image_file:
+                    image_data = image_file.read()
+                    base64_encoded = base64.b64encode(image_data).decode("utf-8")
+                    encoded_images.append(f"data:{context['type']};base64,{base64_encoded}")
+            except (FileNotFoundError, IOError) as e:
+                print(f"Error reading image file {context['value']}: {e}")
+                continue
+
+    return encoded_images
+
+
 def create_ai_optimized_message(
     text: str,
     base64EncodedActiveCellOutput: Optional[str] = None,
-    base64EncodedUploadedImage: Optional[str] = None,
+    additional_context: Optional[List[Dict[str, str]]] = None,
 ) -> ChatCompletionMessageParam:
 
     message_content: Union[str, List[Dict[str, Any]]]
-    has_uploaded_image = (
-        base64EncodedUploadedImage is not None and base64EncodedUploadedImage != ""
+    encoded_images = extract_and_encode_images_from_additional_context(
+        additional_context
     )
+
+    has_uploaded_image = len(encoded_images) > 0
     has_active_cell_output = (
         base64EncodedActiveCellOutput is not None
         and base64EncodedActiveCellOutput != ""
@@ -98,12 +120,12 @@ def create_ai_optimized_message(
             }
         ]
 
-        if has_uploaded_image:
+        for img in encoded_images:
             message_content.append(
                 {
                     "type": "image_url",
                     "image_url": {
-                        "url": f"data:image/png;base64,{base64EncodedUploadedImage}"
+                        "url": img
                     },
                 }
             )

{mito_ai-0.1.42 → mito_ai-0.1.44}/mito_ai/completions/handlers.py

@@ -67,6 +67,7 @@ class CompletionHandler(JupyterHandler, WebSocketHandler):
         self._llm = llm
         self.is_pro = is_pro()
         self._selected_model = FALLBACK_MODEL
+        self.is_electron = False
         identify(llm.key_type)
         
     @property
@@ -128,6 +129,18 @@ class CompletionHandler(JupyterHandler, WebSocketHandler):
             parsed_message = json.loads(message)
             metadata_dict = parsed_message.get('metadata', {})
             type: MessageType = MessageType(parsed_message.get('type'))
+            
+            # Extract environment information from the message
+            environment = parsed_message.get('environment', {})
+            if environment:
+                is_electron = environment.get('isElectron', None)
+                if is_electron is not None:
+                    if is_electron != self.is_electron:
+                        # If the is_electron status is different, log it
+                        identify(key_type=self._llm.key_type, is_electron=is_electron)
+                    
+                    self.is_electron = is_electron
+
         except ValueError as e:
             self.log.error("Invalid completion request.", exc_info=e)
             return

{mito_ai-0.1.42 → mito_ai-0.1.44}/mito_ai/completions/models.py

@@ -3,7 +3,7 @@
 
 import traceback
 from dataclasses import dataclass, field
-from typing import List, Literal, Optional, NewType, Dict
+from typing import List, Literal, Optional, NewType, Dict, Any
 from openai.types.chat import ChatCompletionMessageParam
 from enum import Enum
 from pydantic import BaseModel
@@ -83,7 +83,6 @@ class ChatMessageMetadata():
     variables: Optional[List[str]] = None
     files: Optional[List[str]] = None
     base64EncodedActiveCellOutput: Optional[str] = None
-    base64EncodedUploadedImage: Optional[str] = None
     index: Optional[int] = None
     stream: bool = False
     additionalContext: Optional[List[Dict[str, str]]] = None
@@ -97,7 +96,6 @@ class AgentExecutionMetadata():
     aiOptimizedCells: List[AIOptimizedCell]
     isChromeBrowser: bool
     base64EncodedActiveCellOutput: Optional[str] = None
-    base64EncodedUploadedImage: Optional[str] = None
     variables: Optional[List[str]] = None
     files: Optional[List[str]] = None
     index: Optional[int] = None
@@ -157,6 +155,9 @@ class CompletionRequest:
     # Whether to stream the response (if supported by the model).
     stream: bool = False
     
+    # Environment information from the client
+    environment: Optional[Dict[str, Any]] = None
+    
     
 @dataclass(frozen=True)
 class AICapabilities:

{mito_ai-0.1.42 → mito_ai-0.1.44}/mito_ai/completions/prompt_builders/agent_system_message.py

@@ -51,7 +51,8 @@ Format:
         code: str
         code_summary: str
         cell_type: 'code' | 'markdown'
-    }}
+    }},
+    analysis_assumptions: Optional[List[str]]
 }}
 
 Important information:
@@ -60,7 +61,7 @@ Important information:
 3. The code should be the full contents of that updated code cell. The code that you return will overwrite the existing contents of the code cell so it must contain all necessary code.
 4. The code_summary must be a very short phrase (1–5 words maximum) that begins with a verb ending in "-ing" (e.g., "Loading data", "Filtering rows", "Calculating average", "Plotting revenue"). Avoid full sentences or explanations—this should read like a quick commit message or code label, not a description.
 5. Important: Only use the CELL_UPDATE tool if you want to add/modify a notebook cell in response to the user's request. If the user is just sending you a friendly greeting or asking you a question about yourself, you SHOULD NOT USE A CELL_UPDATE tool because it does not require modifying the notebook. Instead, just use the FINISHED_TASK response.
-6. The assumptions is an optional list of critical assumptions that you made about the data or analysis approach. The assumptions you list here will be displayed to the user so that they can confirm or correct the assumptions. For example: ["NaN values in the impressions column represent 0 impressions", "Only crashes with pedestrian or cyclist fatalities are considered fatal crashes", "Intervention priority combines both volume and severity to identify maximum impact opportunities"].
+6. The analysis_assumptions is an optional list of critical assumptions that you made about the data or analysis approach. The assumptions you list here will be displayed to the user so that they can confirm or correct the assumptions. For example: ["NaN values in the impressions column represent 0 impressions", "Only crashes with pedestrian or cyclist fatalities are considered fatal crashes", "Intervention priority combines both volume and severity to identify maximum impact opportunities"].
 7. Only include important data and analytical assumptions that if incorrect would fundamentally change your analysis conclusions. These should be data handling decisions, methodological choices, and definitional boundaries. Do not include: obvious statements ("Each record is counted once"), result interpretation guidance ("Gaps in the plot represent zero values"), display choices ("Data is sorted for clarity"), internal reasoning ("Bar chart is better than line plot"), or environment assumptions ("Library X is installed"). Prioritize quality over quantity - include only the most critical assumptions or omit the field entirely if there are no critical assumptions made in this step that have not already be shared with the user. If you ever doubt whether an assumption is critical enough to be shared with the user as an assumption, don't include it. Most messages should not include an assumption. 
 8. Do not include the same assumption or variations of the same assumption multiple times in the same conversation. Once you have presented the assumption to the user, they will already have the opportunity to confirm or correct it so do not include it again.
 
@@ -77,7 +78,8 @@ Format:
         code: str   
         code_summary: str
         cell_type: 'code' | 'markdown'
-    }}
+    }},
+    analysis_assumptions: Optional[List[str]]
 }}
 
 Important information:
@@ -86,7 +88,7 @@ Important information:
 3. The code should be the full contents of that updated code cell. The code that you return will overwrite the existing contents of the code cell so it must contain all necessary code.
 4. code_summary must be a very short phrase (1–5 words maximum) that begins with a verb ending in "-ing" (e.g., "Loading data", "Filtering rows", "Calculating average", "Plotting revenue"). Avoid full sentences or explanations—this should read like a quick commit message or code label, not a description.
 5. The cell_type should only be 'markdown' if there is no code to add. There may be times where the code has comments. These are still code cells and should have the cell_type 'code'. Any cells that are labeled 'markdown' will be converted to markdown cells by the user.
-6. The assumptions is an optional list of critical assumptions that you made about the data or analysis approach. The assumptions you list here will be displayed to the user so that they can confirm or correct the assumptions. For example: ["NaN values in the impressions column represent 0 impressions", "Only crashes with pedestrian or cyclist fatalities are considered fatal crashes", "Intervention priority combines both volume and severity to identify maximum impact opportunities"].
+6. The analysis_assumptions is an optional list of critical assumptions that you made about the data or analysis approach. The assumptions you list here will be displayed to the user so that they can confirm or correct the assumptions. For example: ["NaN values in the impressions column represent 0 impressions", "Only crashes with pedestrian or cyclist fatalities are considered fatal crashes", "Intervention priority combines both volume and severity to identify maximum impact opportunities"].
 7. Only include important data and analytical assumptions that if incorrect would fundamentally change your analysis conclusions. These should be data handling decisions, methodological choices, and definitional boundaries. Do not include: obvious statements ("Each record is counted once"), result interpretation guidance ("Gaps in the plot represent zero values"), display choices ("Data is sorted for clarity"), internal reasoning ("Bar chart is better than line plot"), or environment assumptions ("Library X is installed"). Prioritize quality over quantity - include only the most critical assumptions or omit the field entirely if there are no critical assumptions made in this step that have not already be shared with the user. If you ever doubt whether an assumption is critical enough to be shared with the user as an assumption, don't include it. Most messages should not include an assumption. 
 8. Do not include the same assumption or variations of the same assumption multiple times in the same conversation. Once you have presented the assumption to the user, they will already have the opportunity to confirm or correct it so do not include it again.
 

{mito_ai-0.1.42 → mito_ai-0.1.44}/mito_ai/completions/prompt_builders/utils.py

@@ -38,7 +38,7 @@ def get_selected_context_str(additional_context: Optional[List[Dict[str, str]]])
     selected_variables = [context["value"] for context in additional_context if context.get("type") == "variable"]
     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") == "img"]
+    selected_images = [context["value"] for context in additional_context if context.get("type", "").startswith("image/")]
 
     # STEP 2: Create a list of strings (instructions) for each context type
     context_parts = []
@@ -68,5 +68,4 @@ def get_selected_context_str(additional_context: Optional[List[Dict[str, str]]])
         )
 
     # STEP 3: Combine into a single string
-
     return "\n\n".join(context_parts)