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

mito_ai/__init__.py

@@ -14,13 +14,18 @@ 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.file_uploads.urls import get_file_uploads_urls
 
-# 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
+# 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
 import os
-os.environ['MPLBACKEND'] = 'inline'
+os.environ["MPLBACKEND"] = "module://matplotlib_inline.backend_inline"
 
 try:
     from _version import __version__
@@ -87,6 +92,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.41'

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/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]

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/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.
+

mito_ai/file_uploads/handlers.py

@@ -0,0 +1,225 @@
+# Copyright (c) Saga Inc.
+# Distributed under the terms of the GNU Affero General Public License v3.0 License.
+
+import os
+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
+
+
+class FileUploadHandler(APIHandler):
+    # Class-level dictionary to store temporary directories for each file upload
+    # This persists across handler instances since Tornado recreates handlers per request
+    # Key: filename, Value: dict with temp_dir, total_chunks, received_chunks, logged_upload
+    _temp_dirs: Dict[str, Dict[str, Any]] = {}
+
+    def __init__(self, *args: Any, **kwargs: Any) -> None:
+        super().__init__(*args, **kwargs)
+
+    @tornado.web.authenticated
+    def post(self) -> None:
+        """Handle file upload with multipart form data."""
+        try:
+            # Validate request has file
+            if not self._validate_file_upload():
+                return
+
+            uploaded_file = self.request.files["file"][0]
+            filename = uploaded_file["filename"]
+            file_data = uploaded_file["body"]
+
+            # Get notebook directory from request
+            notebook_dir = self.get_argument("notebook_dir", ".")
+
+            # Check if this is a chunked upload
+            chunk_number = self.get_argument("chunk_number", None)
+            total_chunks = self.get_argument("total_chunks", None)
+
+            if chunk_number and total_chunks:
+                self._handle_chunked_upload(
+                    filename, file_data, chunk_number, total_chunks, notebook_dir
+                )
+            else:
+                # Log the file upload attempt for regular (non-chunked) uploads
+                file_extension = filename.split(".")[-1].lower()
+                log_file_upload_attempt(filename, file_extension, False, 0)
+                self._handle_regular_upload(filename, file_data, notebook_dir)
+
+            self.finish()
+
+        except Exception as e:
+            self._handle_error(f"Failed to save file: {str(e)}")
+
+    def _validate_file_upload(self) -> bool:
+        """Validate that a file was uploaded in the request."""
+        if "file" not in self.request.files:
+            self._handle_error("No file uploaded", status_code=400)
+            return False
+        return True
+
+    def _handle_chunked_upload(
+        self,
+        filename: str,
+        file_data: bytes,
+        chunk_number: str,
+        total_chunks: str,
+        notebook_dir: str,
+    ) -> None:
+        """Handle chunked file upload."""
+        chunk_num = int(chunk_number)
+        total_chunks_num = int(total_chunks)
+
+        # Log the file upload attempt only for the first chunk
+        if chunk_num == 1:
+            file_extension = filename.split(".")[-1].lower()
+            log_file_upload_attempt(filename, file_extension, True, total_chunks_num)
+
+        # Save chunk to temporary file
+        self._save_chunk(filename, file_data, chunk_num, total_chunks_num)
+
+        # Check if all chunks are received and reconstruct if complete
+        if self._are_all_chunks_received(filename, total_chunks_num):
+            self._reconstruct_file(filename, total_chunks_num, notebook_dir)
+            self._send_chunk_complete_response(filename, notebook_dir)
+        else:
+            self._send_chunk_received_response(chunk_num, total_chunks_num)
+
+    def _handle_regular_upload(
+        self, filename: str, file_data: bytes, notebook_dir: str
+    ) -> None:
+        """Handle regular (non-chunked) file upload."""
+        file_path = os.path.join(notebook_dir, filename)
+        with open(file_path, "wb") as f:
+            f.write(file_data)
+
+        self.write({"success": True, "filename": filename, "path": file_path})
+
+    def _save_chunk(
+        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}_")
+            self._temp_dirs[filename] = {
+                "temp_dir": temp_dir,
+                "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(
+            self._temp_dirs[filename]["temp_dir"], f"chunk_{chunk_number}"
+        )
+        with open(chunk_filename, "wb") as f:
+            f.write(file_data)
+
+        # 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}")
+
+        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
+            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}")
+        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:
+        """Clean up the temporary directory for a file."""
+        if filename in self._temp_dirs:
+            temp_dir = self._temp_dirs[filename]["temp_dir"]
+            try:
+                import shutil
+
+                shutil.rmtree(temp_dir)
+            except Exception as e:
+                # Log the error but don't fail the upload
+                print(
+                    f"Warning: Failed to clean up temporary directory {temp_dir}: {e}"
+                )
+            finally:
+                # Remove from tracking dictionary
+                del self._temp_dirs[filename]
+
+    def _send_chunk_complete_response(self, filename: str, notebook_dir: str) -> None:
+        """Send response indicating all chunks have been processed and file is complete."""
+        file_path = os.path.join(notebook_dir, filename)
+        self.write(
+            {
+                "success": True,
+                "filename": filename,
+                "path": file_path,
+                "chunk_complete": True,
+            }
+        )
+
+    def _send_chunk_received_response(
+        self, chunk_number: int, total_chunks: int
+    ) -> None:
+        """Send response indicating a chunk was received but file is not yet complete."""
+        self.write(
+            {
+                "success": True,
+                "chunk_received": True,
+                "chunk_number": chunk_number,
+                "total_chunks": total_chunks,
+            }
+        )
+
+    def _handle_error(self, error_message: str, status_code: int = 500) -> None:
+        """Handle errors and send appropriate error response."""
+        log_file_upload_failure(error_message)
+        self.set_status(status_code)
+        self.write({"error": error_message})
+        self.finish()
+
+    def on_finish(self) -> None:
+        """Clean up any remaining temporary directories when the handler is finished."""
+        super().on_finish()
+        # Note: We don't clean up here anymore since we want to preserve state across requests
+        # The cleanup happens when the file is fully reconstructed

mito_ai/file_uploads/urls.py

@@ -0,0 +1,21 @@
+# Copyright (c) Saga Inc.
+# Distributed under the terms of the GNU Affero General Public License v3.0 License.
+
+from typing import List, Tuple, Any
+from jupyter_server.utils import url_path_join
+from mito_ai.file_uploads.handlers import FileUploadHandler
+
+
+def get_file_uploads_urls(base_url: str) -> List[Tuple[str, Any, dict]]:
+    """Get all file uploads related URL patterns.
+
+    Args:
+        base_url: The base URL for the Jupyter server
+
+    Returns:
+        List of (url_pattern, handler_class, handler_kwargs) tuples
+    """
+    BASE_URL = base_url + "/mito-ai"
+    return [
+        (url_path_join(BASE_URL, "upload"), FileUploadHandler, {}),
+    ]

mito_ai/openai_client.py

@@ -35,7 +35,7 @@ from mito_ai.utils.telemetry_utils import (
     USER_KEY,
 )
 
-OPENAI_MODEL_FALLBACK = "gpt-5"
+OPENAI_MODEL_FALLBACK = "gpt-4.1"
 
 class OpenAIClient(LoggingConfigurable):
     """Provide AI feature through OpenAI services."""

mito_ai/tests/file_uploads/__init__.py

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