mito-ai 0.1.40__py3-none-any.whl → 0.1.42__py3-none-any.whl

mito_ai/__init__.py

@@ -14,13 +14,20 @@ from mito_ai.settings.urls import get_settings_urls
 from mito_ai.rules.urls import get_rules_urls
 from mito_ai.auth.urls import get_auth_urls
 from mito_ai.streamlit_preview.urls import get_streamlit_preview_urls
+from mito_ai.app_manager.handlers import AppManagerHandler
+from mito_ai.file_uploads.urls import get_file_uploads_urls
+
+# Force Matplotlib to use the Jupyter inline backend.
+# Background: importing Streamlit sets os.environ["MPLBACKEND"] = "Agg" very early.
+# In a Jupyter kernel, that selects a non‑interactive canvas and can trigger:
+#   "UserWarning: FigureCanvasAgg is non-interactive, and thus cannot be shown"
+# which prevents figures from rendering in notebook outputs.
+# We preempt this by selecting the canonical Jupyter inline backend BEFORE any
+# Matplotlib import, so figures render inline reliably. This must run very early.
+# See: https://github.com/streamlit/streamlit/issues/9640
 
-# Sometimes matplotlib figures do not show up in the notebook with this warning: 
-# UserWarning: FigureCanvasAgg is non-interactive, and thus cannot be shown
-# I believe that streamlit is reconfiguring the matplotlib settings and this is happening as a result.
-# For now, we just set the backend to inline, so that the figures show up again
 import os
-os.environ['MPLBACKEND'] = 'inline'
+os.environ["MPLBACKEND"] = "module://matplotlib_inline.backend_inline"
 
 try:
     from _version import __version__
@@ -77,6 +84,11 @@ def _load_jupyter_server_extension(server_app) -> None: # type: ignore
             url_path_join(base_url, "mito-ai", "version-check"),
             VersionCheckHandler,
             {},
+        ),
+        (
+            url_path_join(base_url, "mito-ai", "app-manager"),
+            AppManagerHandler,
+            {}
         )
     ]
     
@@ -87,6 +99,7 @@ def _load_jupyter_server_extension(server_app) -> None: # type: ignore
     handlers.extend(get_log_urls(base_url, open_ai_provider.key_type))  # type: ignore
     handlers.extend(get_auth_urls(base_url))  # type: ignore
     handlers.extend(get_streamlit_preview_urls(base_url))  # type: ignore
-    
+    handlers.extend(get_file_uploads_urls(base_url)) # type: ignore
+
     web_app.add_handlers(host_pattern, handlers)
     server_app.log.info("Loaded the mito_ai server extension")

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.40'
+__version__ = VERSION = '0.1.42'

mito_ai/app_builder/handlers.py

@@ -156,9 +156,8 @@ class AppBuilderHandler(BaseWebSocketHandler):
                 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)
-                app_path = app_path_result
 
-            deploy_url = await self._deploy_app(app_path, jwt_token)
+            deploy_url = await self._deploy_app(app_directory, jwt_token)
 
             # Send the response
             self.reply(BuildAppReply(

mito_ai/app_manager/__init__.py

@@ -0,0 +1,4 @@
+# Copyright (c) Saga Inc.
+# Distributed under the terms of the GNU Affero General Public License v3.0 License.
+
+"""App manager module for Mito AI.""" 

mito_ai/app_manager/handlers.py

@@ -0,0 +1,129 @@
+# Copyright (c) Saga Inc.
+# Distributed under the terms of the GNU Affero General Public License v3.0 License.
+
+# app_manager/handlers.py
+import os
+import time
+import logging
+from typing import Union
+from mito_ai.utils.websocket_base import BaseWebSocketHandler
+from mito_ai.app_manager.models import (
+    App,
+    AppManagerError,
+    ManageAppRequest,
+    ManageAppReply,
+    ErrorMessage,
+    MessageType
+)
+from mito_ai.constants import ACTIVE_STREAMLIT_BASE_URL
+from mito_ai.logger import get_logger
+import requests
+
+
+class AppManagerHandler(BaseWebSocketHandler):
+    """Handler for app management requests."""
+
+    def initialize(self) -> None:
+        """Initialize the WebSocket handler."""
+        super().initialize()
+        self.log.debug("Initializing app manager websocket connection %s", self.request.path)
+
+    @property
+    def log(self) -> logging.Logger:
+        """Use Mito AI logger."""
+        return get_logger()
+
+    async def on_message(self, message: Union[str, bytes]) -> None:
+        """Handle incoming messages on the WebSocket."""
+        start = time.time()
+
+        # Convert bytes to string if needed
+        if isinstance(message, bytes):
+            message = message.decode('utf-8')
+
+        self.log.debug("App manager 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')
+            message_id = parsed_message.get('message_id')
+
+            if message_type == MessageType.MANAGE_APP.value:
+                # Handle manage app request
+                manage_app_request = ManageAppRequest(**parsed_message)
+                await self._handle_manage_app(manage_app_request)
+            else:
+                self.log.error(f"Unknown message type: {message_type}")
+                error_response = ErrorMessage(
+                    error_type="InvalidRequest",
+                    title=f"Unknown message type: {message_type}",
+                    message_id=message_id
+                )
+                self.reply(error_response)
+
+        except ValueError as e:
+            self.log.error("Invalid app manager request", exc_info=e)
+            error_response = ErrorMessage(
+                error_type=type(e).__name__,
+                title=str(e),
+                message_id=parsed_message.get('message_id') if 'parsed_message' in locals() else None
+            )
+            self.reply(error_response)
+        except Exception as e:
+            self.log.error("Error handling app manager message", exc_info=e)
+            error_response = ErrorMessage(
+                error_type=type(e).__name__,
+                title=str(e),
+                message_id=parsed_message.get('message_id') if 'parsed_message' in locals() else None
+            )
+            self.reply(error_response)
+
+        latency_ms = round((time.time() - start) * 1000)
+        self.log.info(f"App manager handler processed in {latency_ms} ms.")
+
+    async def _handle_manage_app(self, request: ManageAppRequest) -> None:
+        """Handle a manage app request with hardcoded data."""
+        try:
+            jwt_token = request.jwt_token
+            headers = {}
+            if jwt_token and jwt_token != 'placeholder-jwt-token':
+                headers['Authorization'] = f'Bearer {jwt_token}'
+            else:
+                self.log.warning("No JWT token provided for API request")
+                return
+
+            manage_apps_response = requests.get(f"{ACTIVE_STREAMLIT_BASE_URL}/manage-apps",
+                                        headers=headers)
+            manage_apps_response.raise_for_status()
+
+            apps_data = manage_apps_response.json()
+
+            # Create successful response
+            reply = ManageAppReply(
+                apps=apps_data,
+                message_id=request.message_id
+            )
+            self.reply(reply)
+
+        except Exception as e:
+            self.log.error(f"Error handling manage app request: {e}", exc_info=e)
+            
+            try:
+                error = AppManagerError.from_exception(e)
+            except Exception:
+                error = AppManagerError(
+                    error_type=type(e).__name__,
+                    title=str(e)
+                )
+
+            # Return error response
+            error_reply = ManageAppReply(
+                apps=[],
+                error=error,
+                message_id=request.message_id
+            )
+            self.reply(error_reply)

mito_ai/app_manager/models.py

@@ -0,0 +1,58 @@
+# Copyright (c) Saga Inc.
+
+# Distributed under the terms of the GNU Affero General Public License v3.0 License.
+
+from dataclasses import dataclass, field
+from enum import Enum
+from typing import List, Optional
+
+class MessageType(str, Enum):
+    """Types of app manager messages."""
+    MANAGE_APP = "manage-app"
+
+
+@dataclass(frozen=True)
+class ManageAppRequest:
+    """Request to manage apps."""
+    type: str = "manage-app"
+    jwt_token: Optional[str] = None
+    message_id: Optional[str] = None  
+
+@dataclass(frozen=True)
+class App:
+    """App information."""
+    app_name: str
+    url: str
+    status: str
+    created_at: str
+
+@dataclass(frozen=True)
+class AppManagerError:
+    """Error information for app manager operations."""
+    error_type: str
+    title: str
+    traceback: Optional[str] = None
+
+    @classmethod
+    def from_exception(cls, exc: Exception) -> 'AppManagerError':
+        return cls(
+            error_type=type(exc).__name__,
+            title=str(exc),
+            traceback=str(exc)
+        )
+
+@dataclass(frozen=True)
+class ManageAppReply:
+    """Reply to a manage app request."""
+    type: str = "manage-app"
+    apps: List[App] = field(default_factory=list)
+    error: Optional[AppManagerError] = None
+    message_id: Optional[str] = None  
+
+@dataclass(frozen=True)
+class ErrorMessage:
+    """Error message."""
+    error_type: str
+    title: str
+    traceback: Optional[str] = None
+    message_id: Optional[str] = None 

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)
+        new_ai_optimized_message = create_ai_optimized_message(prompt, metadata.base64EncodedActiveCellOutput, metadata.base64EncodedUploadedImage)
         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/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)
+        new_ai_optimized_message = create_ai_optimized_message(prompt, metadata.base64EncodedActiveCellOutput, metadata.base64EncodedUploadedImage)
         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)
+        new_ai_optimized_message = create_ai_optimized_message(prompt, metadata.base64EncodedActiveCellOutput, metadata.base64EncodedUploadedImage)
         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/completions/completion_handlers/utils.py

@@ -6,25 +6,33 @@ from mito_ai.completions.message_history import GlobalMessageHistory
 from mito_ai.completions.models import ThreadID
 from mito_ai.completions.providers import OpenAIProvider
 from openai.types.chat import ChatCompletionMessageParam
-from mito_ai.completions.prompt_builders.chat_system_message import create_chat_system_message_prompt
-from mito_ai.completions.prompt_builders.agent_system_message import create_agent_system_message_prompt
+from mito_ai.completions.prompt_builders.chat_system_message import (
+    create_chat_system_message_prompt,
+)
+from mito_ai.completions.prompt_builders.agent_system_message import (
+    create_agent_system_message_prompt,
+)
+
 
 async def append_chat_system_message(
-        message_history: GlobalMessageHistory,
-        model: str,
-        provider: OpenAIProvider,
-        thread_id: ThreadID
+    message_history: GlobalMessageHistory,
+    model: str,
+    provider: OpenAIProvider,
+    thread_id: ThreadID,
 ) -> None:
-    
+
     # If the system message already exists, do nothing
-    if any(msg["role"] == "system" for msg in message_history.get_ai_optimized_history(thread_id)):
+    if any(
+        msg["role"] == "system"
+        for msg in message_history.get_ai_optimized_history(thread_id)
+    ):
         return
-    
+
     system_message_prompt = create_chat_system_message_prompt()
 
     system_message: ChatCompletionMessageParam = {
         "role": "system",
-        "content": system_message_prompt
+        "content": system_message_prompt,
     }
 
     await message_history.append_message(
@@ -32,54 +40,86 @@ async def append_chat_system_message(
         display_message=system_message,
         model=model,
         llm_provider=provider,
-        thread_id=thread_id
+        thread_id=thread_id,
     )
 
+
 async def append_agent_system_message(
-        message_history: GlobalMessageHistory,
-        model: str,
-        provider: OpenAIProvider,
-        thread_id: ThreadID,
-        isChromeBrowser: bool
+    message_history: GlobalMessageHistory,
+    model: str,
+    provider: OpenAIProvider,
+    thread_id: ThreadID,
+    isChromeBrowser: bool,
 ) -> None:
-    
+
     # If the system message already exists, do nothing
-    if any(msg["role"] == "system" for msg in message_history.get_ai_optimized_history(thread_id)):
+    if any(
+        msg["role"] == "system"
+        for msg in message_history.get_ai_optimized_history(thread_id)
+    ):
         return
-    
+
     system_message_prompt = create_agent_system_message_prompt(isChromeBrowser)
-    
+
     system_message: ChatCompletionMessageParam = {
         "role": "system",
-        "content": system_message_prompt
+        "content": system_message_prompt,
     }
-    
+
     await message_history.append_message(
         ai_optimized_message=system_message,
         display_message=system_message,
         model=model,
         llm_provider=provider,
-        thread_id=thread_id
+        thread_id=thread_id,
     )
-    
-def create_ai_optimized_message(text: str, base64EncodedActiveCellOutput: Optional[str] = None) -> ChatCompletionMessageParam:
+
+
+def create_ai_optimized_message(
+    text: str,
+    base64EncodedActiveCellOutput: Optional[str] = None,
+    base64EncodedUploadedImage: Optional[str] = None,
+) -> ChatCompletionMessageParam:
 
     message_content: Union[str, List[Dict[str, Any]]]
-    if base64EncodedActiveCellOutput is not None and base64EncodedActiveCellOutput != '':
-       message_content = [
+    has_uploaded_image = (
+        base64EncodedUploadedImage is not None and base64EncodedUploadedImage != ""
+    )
+    has_active_cell_output = (
+        base64EncodedActiveCellOutput is not None
+        and base64EncodedActiveCellOutput != ""
+    )
+
+    if has_uploaded_image or has_active_cell_output:
+        message_content = [
             {
                 "type": "text",
                 "text": text,
-            },
-            {
-                "type": "image_url",
-                "image_url": {"url": f"data:image/png;base64,{base64EncodedActiveCellOutput}"},
             }
-       ]
+        ]
+
+        if has_uploaded_image:
+            message_content.append(
+                {
+                    "type": "image_url",
+                    "image_url": {
+                        "url": f"data:image/png;base64,{base64EncodedUploadedImage}"
+                    },
+                }
+            )
+
+        if has_active_cell_output:
+            message_content.append(
+                {
+                    "type": "image_url",
+                    "image_url": {
+                        "url": f"data:image/png;base64,{base64EncodedActiveCellOutput}"
+                    },
+                }
+            )
     else:
         message_content = text
-        
-    return cast(ChatCompletionMessageParam, {
-        "role": "user",
-        "content": message_content
-    })
+
+    return cast(
+        ChatCompletionMessageParam, {"role": "user", "content": message_content}
+    )

mito_ai/completions/handlers.py

@@ -46,7 +46,7 @@ from mito_ai.completions.completion_handlers.agent_execution_handler import get_
 from mito_ai.completions.completion_handlers.agent_auto_error_fixup_handler import get_agent_auto_error_fixup_completion
 from mito_ai.utils.telemetry_utils import identify
 
-FALLBACK_MODEL = "gpt-5"  # Default model to use for safety
+FALLBACK_MODEL = "gpt-4.1"  # Default model to use for safety
 
 # The GlobalMessageHistory is responsible for updating the message histories stored in the .mito/ai-chats directory.
 # We create one GlobalMessageHistory per backend server instance instead of one per websocket connection so that the

mito_ai/completions/message_history.py

@@ -251,7 +251,15 @@ class GlobalMessageHistory:
         with self._lock:
             if thread_id not in self._chat_threads:
                 return []
-            return self._chat_threads[thread_id].display_history
+            
+            thread = self._chat_threads[thread_id]
+            display_history = thread.display_history
+            
+            # When we get a thread, update it's last interaction time so that if the 
+            # user refreshes their browser, this chat will re-appear as the last opened chat.
+            self._update_last_interaction(thread)
+            self._save_thread_to_disk(thread)
+            return display_history
 
     async def append_message(
         self, 

mito_ai/completions/models.py

@@ -29,7 +29,7 @@ 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', 'finished_task']
+    type: Literal['cell_update', 'get_cell_output', 'run_all_cells', 'finished_task']
     message: str
     cell_update: Optional[CellUpdate]
     get_cell_output_cell_id: Optional[str]
@@ -83,6 +83,7 @@ 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
@@ -96,6 +97,7 @@ 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

mito_ai/completions/prompt_builders/agent_execution_prompt.py

@@ -13,12 +13,14 @@ from mito_ai.completions.prompt_builders.utils import (
     get_selected_context_str,
 )
 
+
 def create_agent_execution_prompt(md: AgentExecutionMetadata) -> str:
     variables_str = '\n'.join([f"{variable}" for variable in md.variables or []])
     files_str = '\n'.join([f"{file}" for file in md.files or []])
     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)
+        
     context_str = f"""Remember to choose the correct tool to respond with.
 
 {rules_str}

mito_ai/completions/prompt_builders/agent_smart_debug_prompt.py

@@ -54,6 +54,14 @@ ERROR CORRECTION:
 - Reuse as much of the existing code as possible.
 - DO NOT ADD TEMPORARY COMMENTS like '# Fixed the typo here' or '# Added this line to fix the error'
 - If you encounter a ModuleNotFoundError, you can install the package by adding the the following line to the top of the code cell: `!pip install <package_name> --quiet`.
+- If you encounter a NameError, you can use the RUN_ALL_CELLS tool to run all cells from the top of the notebook to the bottom to bring the variable into scope.
+    RUN_ALL_CELLS:
+    When you want to execute all cells in the notebook from top to bottom, respond with this format:
+    {{
+        type: 'run_all_cells',
+        message: str
+    }}
+    Note that if the name error persists even after using run_all_cells, it means that the variable is not defined in the notebook and you should not reuse this tool. Additionally, this tool could also be used to refresh the notebook state.
 
 <Example>
 

mito_ai/completions/prompt_builders/agent_system_message.py

@@ -206,6 +206,23 @@ Important information:
 ===='''
 }
 
+TOOL: RUN_ALL_CELLS
+
+When you want to execute all cells in the notebook from top to bottom, respond with this format:
+
+{{
+    type: 'run_all_cells',
+    message: str
+}}
+
+Important information:
+1. Use this tool when you encounter a NameError. For example, if you get an error like "NameError: name 'prompts_df' is not defined", you should use this tool to run all cells from the top of the notebook to the bottom to bring the variable into scope.
+2. Note that if the name error persists even after using run_all_cells, it means that the variable is not defined in the notebook and you should not reuse this tool.
+3. Additionally, this tool could also be used to refresh the notebook state.
+4. If running all cells results in an error, the system will automatically handle the error through the normal error fixing process.
+5. Do not use this tool repeatedly if it continues to produce errors - instead, focus on fixing the specific error that occurred.
+====
+
 TOOL: FINISHED_TASK
 
 When you have completed the user's task, respond with a message in this format:

mito_ai/completions/prompt_builders/utils.py

@@ -38,6 +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"]
 
     # STEP 2: Create a list of strings (instructions) for each context type
     context_parts = []
@@ -60,6 +61,12 @@ def get_selected_context_str(additional_context: Optional[List[Dict[str, str]]])
             + "\n".join(selected_db_connections)
         )
 
+    if len(selected_images) > 0:
+        context_parts.append(
+            "The following images have been selected by the user to be used in the task:\n"
+            + "\n".join(selected_images)
+        )
+
     # STEP 3: Combine into a single string
 
     return "\n\n".join(context_parts)

mito_ai/constants.py

@@ -24,8 +24,9 @@ AZURE_OPENAI_ENDPOINT = os.environ.get("AZURE_OPENAI_ENDPOINT")
 AZURE_OPENAI_MODEL = os.environ.get("AZURE_OPENAI_MODEL")
 
 # Mito AI Base URLs and Endpoint Paths
-MITO_PROD_BASE_URL = "https://7eax4i53f5odkshhlry4gw23by0yvnuv.lambda-url.us-east-1.on.aws/v1"
-MITO_DEV_BASE_URL = "https://g5vwmogjg7gh7aktqezyrvcq6a0hyfnr.lambda-url.us-east-1.on.aws/v1"
+MITO_PROD_BASE_URL = "https://7eax4i53f5odkshhlry4gw23by0yvnuv.lambda-url.us-east-1.on.aws/v2"
+MITO_DEV_BASE_URL = "https://g5vwmogjg7gh7aktqezyrvcq6a0hyfnr.lambda-url.us-east-1.on.aws/v2"
+MITO_LOCAL_BASE_URL = "http://127.0.0.1:8000/v2" # When you are running the mito completion server locally
 
 # Set ACTIVE_BASE_URL manually
 ACTIVE_BASE_URL = MITO_PROD_BASE_URL  # Change to MITO_DEV_BASE_URL for dev

mito_ai/file_uploads/__init__.py

@@ -0,0 +1,3 @@
+# Copyright (c) Saga Inc.
+# Distributed under the terms of the GNU Affero General Public License v3.0 License.
+