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

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

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.

mito_ai/tests/file_uploads/test_handlers.py

@@ -0,0 +1,267 @@
+# Copyright (c) Saga Inc.
+# Distributed under the terms of the GNU Affero General Public License v3.0 License.
+
+import os
+import tempfile
+import pytest
+from unittest.mock import Mock, patch
+import tornado.web
+from tornado.httputil import HTTPServerRequest
+from tornado.web import Application
+
+from mito_ai.file_uploads.handlers import FileUploadHandler
+
+
+@pytest.fixture
+def temp_dir():
+    """Create a temporary directory for test files."""
+    temp_dir = tempfile.mkdtemp()
+    original_cwd = os.getcwd()
+    os.chdir(temp_dir)
+    yield temp_dir
+    os.chdir(original_cwd)
+    # Clean up temporary files
+    for file in os.listdir(temp_dir):
+        os.remove(os.path.join(temp_dir, file))
+    os.rmdir(temp_dir)
+
+
+@pytest.fixture
+def handler():
+    """Create a FileUploadHandler instance for testing."""
+    app = Application()
+    request = HTTPServerRequest(method="POST", uri="/upload")
+
+    # Mock the connection to avoid Tornado's assertion
+    request.connection = Mock()
+
+    handler = FileUploadHandler(app, request)
+
+    # Mock methods properly to avoid mypy errors
+    handler.write = Mock()  # type: ignore
+    handler.finish = Mock()  # type: ignore
+    handler.set_status = Mock()  # type: ignore
+    handler.get_argument = Mock()  # type: ignore
+
+    # Mock authentication for Jupyter server
+    handler._jupyter_current_user = "test_user"  # type: ignore
+
+    return handler
+
+
+def test_validate_file_upload_success(handler):
+    """Test successful file upload validation."""
+    handler.request.files = {"file": [Mock(filename="test.csv", body=b"data")]}  # type: ignore
+    result = handler._validate_file_upload()
+    assert result is True
+
+
+def test_validate_file_upload_failure(handler):
+    """Test file upload validation when no file is present."""
+    handler.request.files = {}  # type: ignore
+    result = handler._validate_file_upload()
+    assert result is False
+    handler.set_status.assert_called_with(400)
+
+
+def test_regular_upload_success(handler, temp_dir):
+    """Test successful regular (non-chunked) file upload."""
+    filename = "test.csv"
+    file_data = b"test,data\n1,2"
+    notebook_dir = temp_dir
+
+    handler._handle_regular_upload(filename, file_data, notebook_dir)
+
+    # Verify file was written
+    file_path = os.path.join(notebook_dir, filename)
+    with open(file_path, "rb") as f:
+        content = f.read()
+    assert content == file_data
+
+    # Verify response
+    handler.write.assert_called_with(
+        {"success": True, "filename": filename, "path": file_path}
+    )
+
+
+def test_chunked_upload_first_chunk(handler, temp_dir):
+    """Test handling first chunk of a chunked upload."""
+    filename = "large_file.csv"
+    file_data = b"chunk1_data"
+    chunk_number = "1"
+    total_chunks = "3"
+    notebook_dir = temp_dir
+
+    handler._handle_chunked_upload(
+        filename, file_data, chunk_number, total_chunks, notebook_dir
+    )
+
+    # Verify chunk was saved (check temp dir structure)
+    assert filename in handler._temp_dirs
+    temp_dir_path = handler._temp_dirs[filename]["temp_dir"]
+    chunk_file = os.path.join(temp_dir_path, "chunk_1")
+    assert os.path.exists(chunk_file)
+
+    # Verify response indicates chunk received but not complete
+    handler.write.assert_called_with(
+        {
+            "success": True,
+            "chunk_received": True,
+            "chunk_number": 1,
+            "total_chunks": 3,
+        }
+    )
+
+
+def test_chunked_upload_completion(handler, temp_dir):
+    """Test completing a chunked upload when all chunks are received."""
+    filename = "large_file.csv"
+    total_chunks = 2
+    notebook_dir = temp_dir
+
+    # Process first chunk
+    handler._handle_chunked_upload(
+        filename, b"chunk1_data", "1", str(total_chunks), notebook_dir
+    )
+
+    # Process final chunk
+    handler._handle_chunked_upload(
+        filename, b"chunk2_data", "2", str(total_chunks), notebook_dir
+    )
+
+    # Verify final file was created
+    file_path = os.path.join(notebook_dir, filename)
+    assert os.path.exists(file_path)
+    with open(file_path, "rb") as f:
+        content = f.read()
+    assert content == b"chunk1_datachunk2_data"
+
+    # Verify temp dir was cleaned up
+    assert filename not in handler._temp_dirs
+
+    # Verify completion response
+    handler.write.assert_called_with(
+        {
+            "success": True,
+            "filename": filename,
+            "path": file_path,
+            "chunk_complete": True,
+        }
+    )
+
+
+def test_error_handling(handler):
+    """Test error handling in upload process."""
+    error_message = "Test error message"
+    status_code = 500
+
+    handler._handle_error(error_message, status_code)
+
+    handler.set_status.assert_called_with(status_code)
+    handler.write.assert_called_with({"error": error_message})
+    handler.finish.assert_called_once()
+
+
+@patch("mito_ai.file_uploads.handlers.FileUploadHandler._validate_file_upload")
+def test_post_method_regular_upload(mock_validate, handler):
+    """Test POST method for regular upload."""
+    mock_validate.return_value = True
+    handler.request.files = {"file": [Mock(filename="test.csv", body=b"data")]}  # type: ignore
+    handler.get_argument.return_value = None  # No chunk parameters
+
+    handler.post()
+
+    mock_validate.assert_called_once()
+    handler.finish.assert_called_once()
+
+
+@patch("mito_ai.file_uploads.handlers.FileUploadHandler._validate_file_upload")
+def test_post_method_chunked_upload(mock_validate, handler):
+    """Test POST method for chunked upload."""
+    mock_validate.return_value = True
+    handler.request.files = {"file": [Mock(filename="test.csv", body=b"data")]}  # type: ignore
+    handler.get_argument.side_effect = lambda name, default=None: {
+        "chunk_number": "1",
+        "total_chunks": "3",
+    }.get(name, default)
+
+    handler.post()
+
+    mock_validate.assert_called_once()
+    handler.finish.assert_called_once()
+
+
+def test_are_all_chunks_received_true(handler, temp_dir):
+    """Test that all chunks are detected when present."""
+    filename = "test.csv"
+    total_chunks = 2
+
+    # Manually set up the temp dir structure
+    temp_dir_path = tempfile.mkdtemp(prefix=f"mito_upload_{filename}_")
+    handler._temp_dirs[filename] = {
+        "temp_dir": temp_dir_path,
+        "total_chunks": total_chunks,
+        "received_chunks": {1, 2},
+    }
+
+    result = handler._are_all_chunks_received(filename, total_chunks)
+    assert result is True
+
+    # Clean up
+    import shutil
+
+    shutil.rmtree(temp_dir_path)
+
+
+def test_are_all_chunks_received_false(handler, temp_dir):
+    """Test that missing chunks are detected."""
+    filename = "test.csv"
+    total_chunks = 2
+
+    # Manually set up the temp dir structure with only one chunk
+    temp_dir_path = tempfile.mkdtemp(prefix=f"mito_upload_{filename}_")
+    handler._temp_dirs[filename] = {
+        "temp_dir": temp_dir_path,
+        "total_chunks": total_chunks,
+        "received_chunks": {1},  # Only chunk 1 received
+    }
+
+    result = handler._are_all_chunks_received(filename, total_chunks)
+    assert result is False
+
+    # Clean up
+    import shutil
+
+    shutil.rmtree(temp_dir_path)
+
+
+def test_save_chunk(handler, temp_dir):
+    """Test saving individual chunks."""
+    filename = "test.csv"
+    file_data = b"chunk_data"
+    chunk_number = 1
+    total_chunks = 3
+
+    # Mock the file operations to avoid filesystem issues
+    with patch("builtins.open", create=True) as mock_open:
+        mock_file = Mock()
+        mock_open.return_value.__enter__.return_value = mock_file
+
+        handler._save_chunk(filename, file_data, chunk_number, total_chunks)
+
+        # Verify temp dir was created in the handler's tracking
+        assert filename in handler._temp_dirs
+        temp_dir_path = handler._temp_dirs[filename]["temp_dir"]
+
+        # Verify the expected chunk filename was used
+        expected_chunk_filename = os.path.join(temp_dir_path, f"chunk_{chunk_number}")
+        mock_open.assert_called_with(expected_chunk_filename, "wb")
+
+        # Verify file data was written
+        mock_file.write.assert_called_with(file_data)
+
+        # Verify chunk was marked as received
+        assert chunk_number in handler._temp_dirs[filename]["received_chunks"]
+
+        # Clean up
+        del handler._temp_dirs[filename]

mito_ai/tests/message_history/test_message_history_utils.py

@@ -7,9 +7,9 @@ from openai.types.chat import ChatCompletionMessageParam
 from mito_ai.utils.message_history_utils import trim_sections_from_message_content, trim_old_messages
 from mito_ai.completions.prompt_builders.chat_prompt import create_chat_prompt
 from mito_ai.completions.prompt_builders.agent_execution_prompt import create_agent_execution_prompt
-from mito_ai.completions.prompt_builders.agent_smart_debug_prompt import (
-    create_agent_smart_debug_prompt,
-)
+from mito_ai.completions.prompt_builders.agent_smart_debug_prompt import create_agent_smart_debug_prompt
+from unittest.mock import Mock, patch
+from mito_ai.completions.message_history import GlobalMessageHistory, ChatThread
 from mito_ai.completions.prompt_builders.smart_debug_prompt import create_error_prompt
 from mito_ai.completions.prompt_builders.explain_code_prompt import create_explain_code_prompt
 from mito_ai.completions.models import (
@@ -27,6 +27,9 @@ from mito_ai.completions.prompt_builders.prompt_constants import (
     CONTENT_REMOVED_PLACEHOLDER,
 )
 
+
+
+
 # Standard test data for multiple tests
 TEST_VARIABLES = ["'df': pd.DataFrame({'col1': [1, 2, 3], 'col2': [4, 5, 6]})"]
 TEST_FILES = ["data.csv", "script.py"]
@@ -386,4 +389,54 @@ def test_trim_mixed_content_messages() -> None:
     
     # Verify that the recent messages are untouched
     assert trimmed_messages[1] == message_list[1]
-    assert trimmed_messages[2] == message_list[2]
+    assert trimmed_messages[2] == message_list[2]
+
+
+def test_get_display_history_calls_update_last_interaction() -> None:
+    """Test that get_display_history calls _update_last_interaction when retrieving a thread."""
+    
+    # Create a mock thread
+    thread_id = ThreadID("test-thread-id")
+    mock_thread = Mock(spec=ChatThread)
+    mock_thread.display_history = [{"role": "user", "content": "test message"}]
+    mock_thread.last_interaction_ts = 1234567890.0
+    
+    # Create message history instance and add the mock thread
+    message_history = GlobalMessageHistory()
+    message_history._chat_threads = {thread_id: mock_thread}
+    
+    # Mock the _update_last_interaction method
+    with patch.object(message_history, '_update_last_interaction') as mock_update:
+        with patch.object(message_history, '_save_thread_to_disk') as mock_save:
+            # Call get_display_history
+            result = message_history.get_display_history(thread_id)
+            
+            # Verify _update_last_interaction was called with the thread
+            mock_update.assert_called_once_with(mock_thread)
+            
+            # Verify _save_thread_to_disk was also called
+            mock_save.assert_called_once_with(mock_thread)
+            
+            # Verify the result is correct
+            assert result == [{"role": "user", "content": "test message"}]
+
+
+def test_get_display_history_returns_empty_for_nonexistent_thread() -> None:
+    """Test that get_display_history returns empty list for non-existent thread."""
+    from mito_ai.completions.message_history import GlobalMessageHistory
+    from mito_ai.completions.models import ThreadID
+    
+    message_history = GlobalMessageHistory()
+    thread_id = ThreadID("nonexistent-thread-id")
+    
+    # Mock the methods to ensure they're not called
+    with patch.object(message_history, '_update_last_interaction') as mock_update:
+        with patch.object(message_history, '_save_thread_to_disk') as mock_save:
+            result = message_history.get_display_history(thread_id)
+            
+            # Verify methods were not called since thread doesn't exist
+            mock_update.assert_not_called()
+            mock_save.assert_not_called()
+            
+            # Verify empty result
+            assert result == []

mito_ai/utils/mito_server_utils.py

@@ -11,6 +11,7 @@ from tornado.httpclient import HTTPResponse
 from mito_ai.constants import MITO_GEMINI_URL
 from mito_ai.utils.utils import _create_http_client
 
+MITO_ERROR_MARKER = "MITO_ERROR_MARKER:"
 
 class ProviderCompletionException(Exception):
     """Custom exception for Mito server errors that converts well to CompletionError."""
@@ -179,6 +180,12 @@ async def stream_response_from_mito_server(
                 if chunk_processor:
                     processed_chunk = chunk_processor(chunk)
 
+                # Check if this chunk contains an error marker
+                if processed_chunk.startswith(MITO_ERROR_MARKER):
+                    error_message = processed_chunk[len(MITO_ERROR_MARKER):]
+                    print(f"Detected error in {provider_name} stream: {error_message}")
+                    raise ProviderCompletionException(error_message, provider_name=provider_name)
+
                 if reply_fn is not None and message_id is not None:
                     # Send the chunk directly to the frontend
                     reply_fn(CompletionStreamChunk(

mito_ai/utils/server_limits.py

@@ -28,7 +28,7 @@ free tier, but running AI models is expensive, so we need to limit the usage
 or we will no longer be able to provide this free tier.
 """
 # Monthly chat completions limit for free tier users
-OS_MONTHLY_AI_COMPLETIONS_LIMIT: Final[int] = 50
+OS_MONTHLY_AI_COMPLETIONS_LIMIT: Final[int] = 150
 
 # Monthly autocomplete limit for free tier users
 OS_MONTHLY_AUTOCOMPLETE_LIMIT: Final[int] = 5000