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

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.43'

mito_ai/app_manager/handlers.py

@@ -17,6 +17,7 @@ from mito_ai.app_manager.models import (
 )
 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
 
 
@@ -102,6 +103,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,

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
@@ -24,7 +23,7 @@ class App:
     app_name: str
     url: str
     status: str
-    created_at: str
+    last_deployed_at: str
 
 @dataclass(frozen=True)
 class AppManagerError:

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

@@ -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

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)

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:

mito_ai/tests/completions/completion_handlers_utils_test.py

@@ -1,7 +1,28 @@
 # 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
+import base64
+import os
+import tempfile
+from contextlib import contextmanager
+from mito_ai.completions.completion_handlers.utils import (
+    create_ai_optimized_message,
+    extract_and_encode_images_from_additional_context,
+)
+
+
+@contextmanager
+def temporary_image_file(suffix=".png", content=b"fake_image_data"):
+    """Context manager that creates a temporary image file for testing."""
+    with tempfile.NamedTemporaryFile(suffix=suffix, delete=False) as temp_file:
+        temp_file.write(content)
+        temp_file_path = temp_file.name
+
+    try:
+        yield temp_file_path
+    finally:
+        # Clean up the temporary file
+        os.unlink(temp_file_path)
 
 
 def test_text_only_message():
@@ -14,14 +35,41 @@ def test_text_only_message():
 
 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"
-    )
+    with temporary_image_file() as temp_file_path:
+        result = create_ai_optimized_message(
+            text="Analyze this",
+            additional_context=[{"type": "image/png", "value": temp_file_path}],
+        )
+
+        assert result["role"] == "user"
+        assert isinstance(result["content"], list)
+        assert result["content"][0]["type"] == "text"
+        assert result["content"][1]["type"] == "image_url"
 
-    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_multiple_uploaded_images():
+    """Test scenario where the user uploads multiple images"""
+    with temporary_image_file(suffix=".png", content=b"image1_data") as temp_file1:
+        with temporary_image_file(suffix=".jpg", content=b"image2_data") as temp_file2:
+            result = create_ai_optimized_message(
+                text="Analyze these images",
+                additional_context=[
+                    {"type": "image/png", "value": temp_file1},
+                    {"type": "image/jpeg", "value": temp_file2},
+                ],
+            )
+
+            assert result["role"] == "user"
+            assert isinstance(result["content"], list)
+            assert len(result["content"]) == 3  # text + 2 images
+            assert result["content"][0]["type"] == "text"
+            assert result["content"][0]["text"] == "Analyze these images"
+            assert result["content"][1]["type"] == "image_url"
+            assert result["content"][2]["type"] == "image_url"
+            
+            # Verify the image URLs are properly formatted
+            assert result["content"][1]["image_url"]["url"].startswith("data:image/png;base64,")
+            assert result["content"][2]["image_url"]["url"].startswith("data:image/jpeg;base64,")
 
 
 def test_message_with_active_cell_output():
@@ -38,14 +86,105 @@ def test_message_with_active_cell_output():
 
 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",
+    with temporary_image_file() as temp_file_path:
+        result = create_ai_optimized_message(
+            text="Analyze this",
+            additional_context=[{"type": "image/png", "value": temp_file_path}],
+            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"
+
+
+def test_extract_and_encode_images_from_additional_context_valid_image():
+    """Test extracting and encoding a valid image file"""
+    with temporary_image_file() as temp_file_path:
+        additional_context = [{"type": "image/png", "value": temp_file_path}]
+
+        encoded_images = extract_and_encode_images_from_additional_context(
+            additional_context
+        )
+
+        assert len(encoded_images) == 1
+        assert encoded_images[0].startswith("data:image/png;base64,")
+        # Verify it's valid base64 by checking it can be decoded
+        base64_data = encoded_images[0].split(",")[1]
+        decoded_data = base64.b64decode(base64_data)
+        assert decoded_data == b"fake_image_data"
+
+
+def test_extract_and_encode_images_from_additional_context_multiple_images():
+    """Test extracting and encoding multiple image files"""
+    with temporary_image_file(suffix=".png", content=b"image1_data") as temp_file1:
+        with temporary_image_file(suffix=".jpg", content=b"image2_data") as temp_file2:
+            additional_context = [
+                {"type": "image/png", "value": temp_file1},
+                {"type": "image/jpeg", "value": temp_file2},
+            ]
+
+            encoded_images = extract_and_encode_images_from_additional_context(
+                additional_context
+            )
+
+            assert len(encoded_images) == 2
+            assert encoded_images[0].startswith("data:image/png;base64,")
+            assert encoded_images[1].startswith("data:image/jpeg;base64,")
+
+
+def test_extract_and_encode_images_from_additional_context_invalid_file():
+    """Test handling of invalid/non-existent image files"""
+    additional_context = [{"type": "image/png", "value": "non_existent_file.png"}]
+
+    encoded_images = extract_and_encode_images_from_additional_context(
+        additional_context
     )
 
-    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"
+    assert len(encoded_images) == 0
+
+
+def test_extract_and_encode_images_from_additional_context_non_image_types():
+    """Test that non-image types are ignored"""
+    with temporary_image_file(suffix=".txt", content=b"text_data") as temp_file:
+        additional_context = [
+            {"type": "text/plain", "value": temp_file},
+            {"type": "application/pdf", "value": "document.pdf"},
+        ]
+
+        encoded_images = extract_and_encode_images_from_additional_context(
+            additional_context
+        )
+
+        assert len(encoded_images) == 0
+
+
+def test_extract_and_encode_images_from_additional_context_mixed_types():
+    """Test handling of mixed image and non-image types"""
+    with temporary_image_file() as temp_image_file:
+        additional_context = [
+            {"type": "image/png", "value": temp_image_file},
+            {"type": "text/plain", "value": "document.txt"},
+            {"type": "image/jpeg", "value": "non_existent.jpg"},
+        ]
+
+        encoded_images = extract_and_encode_images_from_additional_context(
+            additional_context
+        )
+
+        # Should only have the valid PNG image
+        assert len(encoded_images) == 1
+        assert encoded_images[0].startswith("data:image/png;base64,")
+
+
+def test_extract_and_encode_images_from_additional_context_empty():
+    """Test handling of empty or None additional_context"""
+    # Test with None
+    encoded_images = extract_and_encode_images_from_additional_context(None)
+    assert len(encoded_images) == 0
+
+    # Test with empty list
+    encoded_images = extract_and_encode_images_from_additional_context([])
+    assert len(encoded_images) == 0

mito_ai/tests/file_uploads/test_handlers.py

@@ -265,3 +265,18 @@ def test_save_chunk(handler, temp_dir):
 
         # Clean up
         del handler._temp_dirs[filename]
+
+
+def test_image_size_limit_exceeded(handler, temp_dir):
+    """Test that image uploads exceeding 3MB are rejected."""
+    filename = "large_image.jpg"
+    # Create 5MB of data (5 * 1024 * 1024 bytes)
+    file_data = b"x" * (5 * 1024 * 1024)
+    notebook_dir = temp_dir
+
+    # The _handle_regular_upload should raise a ValueError for oversized images
+    with pytest.raises(ValueError) as exc_info:
+        handler._handle_regular_upload(filename, file_data, notebook_dir)
+    
+    # Verify the error message mentions the size limit
+    assert "exceeded 3MB limit" in str(exc_info.value)

{mito_ai-0.1.42.data → mito_ai-0.1.43.data}/data/share/jupyter/labextensions/mito_ai/build_log.json

@@ -712,7 +712,7 @@
             "semver": {},
             "vscode-diff": {},
             "mito_ai": {
-              "version": "0.1.42",
+              "version": "0.1.43",
               "singleton": true,
               "import": "/home/runner/work/mito/mito/mito-ai/lib/index.js"
             }

{mito_ai-0.1.42.data → mito_ai-0.1.43.data}/data/share/jupyter/labextensions/mito_ai/package.json

@@ -1,6 +1,6 @@
 {
   "name": "mito_ai",
-  "version": "0.1.42",
+  "version": "0.1.43",
   "description": "AI chat for JupyterLab",
   "keywords": [
     "jupyter",
@@ -140,7 +140,7 @@
     "outputDir": "mito_ai/labextension",
     "schemaDir": "schema",
     "_build": {
-      "load": "static/remoteEntry.c7d9d8635826165de52e.js",
+      "load": "static/remoteEntry.502aef26f0416fab7435.js",
       "extension": "./extension",
       "style": "./style"
     }

{mito_ai-0.1.42.data → mito_ai-0.1.43.data}/data/share/jupyter/labextensions/mito_ai/schemas/mito_ai/package.json.orig

@@ -1,6 +1,6 @@
 {
   "name": "mito_ai",
-  "version": "0.1.42",
+  "version": "0.1.43",
   "description": "AI chat for JupyterLab",
   "keywords": [
     "jupyter",