mito-ai 0.1.41__py3-none-any.whl → 0.1.42__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.42'

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/models.py

@@ -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/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/tests/completions/completion_handlers_utils_test.py

@@ -0,0 +1,51 @@
+# Copyright (c) Saga Inc.
+# Distributed under the terms of the GNU Affero General Public License v3.0 License.
+
+from mito_ai.completions.completion_handlers.utils import create_ai_optimized_message
+
+
+def test_text_only_message():
+    """Test scenario where the user only inputs text"""
+    result = create_ai_optimized_message("Hello world")
+
+    assert result["role"] == "user"
+    assert result["content"] == "Hello world"
+
+
+def test_message_with_uploaded_image():
+    """Test scenario where the user uploads an image"""
+    result = create_ai_optimized_message(
+        text="Analyze this", base64EncodedUploadedImage="image_data"
+    )
+
+    assert result["role"] == "user"
+    assert isinstance(result["content"], list)
+    assert result["content"][0]["type"] == "text"
+    assert result["content"][1]["type"] == "image_url"
+
+
+def test_message_with_active_cell_output():
+    """Test scenario where the active cell has an output"""
+    result = create_ai_optimized_message(
+        text="Analyze this", base64EncodedActiveCellOutput="cell_output_data"
+    )
+
+    assert result["role"] == "user"
+    assert isinstance(result["content"], list)
+    assert result["content"][0]["type"] == "text"
+    assert result["content"][1]["type"] == "image_url"
+
+
+def test_message_with_uploaded_image_and_active_cell_output():
+    """Test scenario where the user uploads an image and the active cell has an output"""
+    result = create_ai_optimized_message(
+        text="Analyze this",
+        base64EncodedUploadedImage="image_data",
+        base64EncodedActiveCellOutput="cell_output_data",
+    )
+
+    assert result["role"] == "user"
+    assert isinstance(result["content"], list)
+    assert result["content"][0]["type"] == "text"
+    assert result["content"][1]["type"] == "image_url"
+    assert result["content"][2]["type"] == "image_url"