mito-ai 0.1.41__py3-none-any.whl → 0.1.43__py3-none-any.whl

mito_ai/__init__.py

@@ -14,6 +14,7 @@ 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.
@@ -24,6 +25,7 @@ from mito_ai.file_uploads.urls import get_file_uploads_urls
 # 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
+
 import os
 os.environ["MPLBACKEND"] = "module://matplotlib_inline.backend_inline"
 
@@ -82,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,
+            {}
         )
     ]
     

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.41'
+__version__ = VERSION = '0.1.43'

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,134 @@
+# 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
+from mito_ai.app_manager.utils import convert_utc_to_local_time
+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()
+
+            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,
+                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,57 @@
+# 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
+    last_deployed_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/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/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.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/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.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)
+        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/completions/completion_handlers/utils.py

@@ -1,30 +1,39 @@
 # 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
 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 +41,107 @@ 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 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,
+    additional_context: Optional[List[Dict[str, str]]] = None,
+) -> ChatCompletionMessageParam:
 
     message_content: Union[str, List[Dict[str, Any]]]
-    if base64EncodedActiveCellOutput is not None and base64EncodedActiveCellOutput != '':
-       message_content = [
+    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 != ""
+    )
+
+    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}"},
             }
-       ]
+        ]
+
+        for img in encoded_images:
+            message_content.append(
+                {
+                    "type": "image_url",
+                    "image_url": {
+                        "url": img
+                    },
+                }
+            )
+
+        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/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", "").startswith("image/")]
 
     # STEP 2: Create a list of strings (instructions) for each context type
     context_parts = []
@@ -60,6 +61,11 @@ def get_selected_context_str(additional_context: Optional[List[Dict[str, str]]])
             + "\n".join(selected_db_connections)
         )
 
-    # STEP 3: Combine into a single string
+    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/file_uploads/handlers.py

@@ -6,7 +6,38 @@ import tempfile
 import tornado
 from typing import Dict, Any
 from jupyter_server.base.handlers import APIHandler
-from mito_ai.utils.telemetry_utils import log_file_upload_attempt, log_file_upload_failure
+from mito_ai.utils.telemetry_utils import (
+    log_file_upload_attempt,
+    log_file_upload_failure,
+)
+
+MAX_IMAGE_SIZE_MB = 3
+
+
+def _is_image_file(filename: str) -> bool:
+    image_extensions = {
+        ".jpg",
+        ".jpeg",
+        ".png",
+        ".gif",
+        ".bmp",
+        ".tiff",
+        ".tif",
+        ".webp",
+        ".svg",
+    }
+    file_extension = os.path.splitext(filename)[1].lower()
+    return file_extension in image_extensions
+
+
+def _check_image_size_limit(file_data: bytes, filename: str) -> None:
+    if not _is_image_file(filename):
+        return
+
+    file_size_mb = len(file_data) / (1024 * 1024)  # Convert bytes to MB
+
+    if file_size_mb > MAX_IMAGE_SIZE_MB:
+        raise ValueError(f"Image exceeded {MAX_IMAGE_SIZE_MB}MB limit.")
 
 
 class FileUploadHandler(APIHandler):
@@ -50,7 +81,7 @@ class FileUploadHandler(APIHandler):
             self.finish()
 
         except Exception as e:
-            self._handle_error(f"Failed to save file: {str(e)}")
+            self._handle_error(str(e))
 
     def _validate_file_upload(self) -> bool:
         """Validate that a file was uploaded in the request."""
@@ -90,6 +121,9 @@ class FileUploadHandler(APIHandler):
         self, filename: str, file_data: bytes, notebook_dir: str
     ) -> None:
         """Handle regular (non-chunked) file upload."""
+        # Check image file size limit before saving
+        _check_image_size_limit(file_data, filename)
+
         file_path = os.path.join(notebook_dir, filename)
         with open(file_path, "wb") as f:
             f.write(file_data)
@@ -100,8 +134,6 @@ class FileUploadHandler(APIHandler):
         self, filename: str, file_data: bytes, chunk_number: int, total_chunks: int
     ) -> None:
         """Save a chunk to a temporary file."""
-        print(f"DEBUG: Saving chunk {chunk_number}/{total_chunks} for file {filename}")
-
         # Initialize temporary directory for this file if it doesn't exist
         if filename not in self._temp_dirs:
             temp_dir = tempfile.mkdtemp(prefix=f"mito_upload_{filename}_")
@@ -110,7 +142,6 @@ class FileUploadHandler(APIHandler):
                 "total_chunks": total_chunks,
                 "received_chunks": set(),
             }
-            print(f"DEBUG: Created temp dir {temp_dir} for file {filename}")
 
         # Save the chunk to the temporary directory
         chunk_filename = os.path.join(
@@ -121,28 +152,20 @@ class FileUploadHandler(APIHandler):
 
         # Mark this chunk as received
         self._temp_dirs[filename]["received_chunks"].add(chunk_number)
-        print(
-            f"DEBUG: Saved chunk {chunk_number}, total received: {len(self._temp_dirs[filename]['received_chunks'])}/{total_chunks}"
-        )
 
     def _are_all_chunks_received(self, filename: str, total_chunks: int) -> bool:
         """Check if all chunks for a file have been received."""
         if filename not in self._temp_dirs:
-            print(f"DEBUG: No temp dir found for {filename}")
             return False
 
         received_chunks = self._temp_dirs[filename]["received_chunks"]
         is_complete = len(received_chunks) == total_chunks
-        print(
-            f"DEBUG: Checking completion for {filename}: {len(received_chunks)}/{total_chunks} chunks received, complete: {is_complete}"
-        )
         return is_complete
 
     def _reconstruct_file(
         self, filename: str, total_chunks: int, notebook_dir: str
     ) -> None:
         """Reconstruct the final file from all chunks and clean up temporary directory."""
-        print(f"DEBUG: Starting reconstruction for {filename}")
 
         if filename not in self._temp_dirs:
             raise ValueError(f"No temporary directory found for file: {filename}")
@@ -150,23 +173,23 @@ class FileUploadHandler(APIHandler):
         temp_dir = self._temp_dirs[filename]["temp_dir"]
         file_path = os.path.join(notebook_dir, filename)
 
-        print(f"DEBUG: Reconstructing from {temp_dir} to {file_path}")
-
         try:
-            # Reconstruct the file from chunks
+            # First, read all chunks to check total file size for images
+            all_file_data = b""
+            for i in range(1, total_chunks + 1):
+                chunk_filename = os.path.join(temp_dir, f"chunk_{i}")
+                with open(chunk_filename, "rb") as chunk_file:
+                    chunk_data = chunk_file.read()
+                    all_file_data += chunk_data
+
+            # Check image file size limit before saving
+            _check_image_size_limit(all_file_data, filename)
+
+            # Write the complete file
             with open(file_path, "wb") as final_file:
-                for i in range(1, total_chunks + 1):
-                    chunk_filename = os.path.join(temp_dir, f"chunk_{i}")
-                    print(f"DEBUG: Reading chunk {i} from {chunk_filename}")
-                    with open(chunk_filename, "rb") as chunk_file:
-                        chunk_data = chunk_file.read()
-                        final_file.write(chunk_data)
-                        print(f"DEBUG: Wrote {len(chunk_data)} bytes from chunk {i}")
-
-            print(f"DEBUG: Successfully reconstructed {filename}")
+                final_file.write(all_file_data)
         finally:
             # Clean up the temporary directory
-            print(f"DEBUG: Cleaning up temp dir for {filename}")
             self._cleanup_temp_dir(filename)
 
     def _cleanup_temp_dir(self, filename: str) -> None: