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

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

mito_ai/utils/telemetry_utils.py

@@ -337,15 +337,6 @@ def log_db_connection_success(connection_type: str, schema: Dict[str, Any]) -> N
         },
     )
 
-def log_ai_completion_retry(key_type: Literal['mito_server_key', 'user_key'], message_type: MessageType, error: BaseException) -> None:
-    log(MITO_AI_COMPLETION_RETRY, params={KEY_TYPE_PARAM: key_type, "message_type": message_type}, key_type=key_type, error=error)
-    
-def log_ai_completion_error(key_type: Literal['mito_server_key', 'user_key'], message_type: MessageType, error: BaseException) -> None:
-    log(MITO_AI_COMPLETION_ERROR, params={KEY_TYPE_PARAM: key_type, "message_type": message_type}, key_type=key_type, error=error)
-    
-def log_mito_server_free_tier_limit_reached(key_type: Literal['mito_server_key', 'user_key'], message_type: MessageType) -> None:
-    log(MITO_SERVER_FREE_TIER_LIMIT_REACHED, params={KEY_TYPE_PARAM: key_type, "message_type": message_type}, key_type=key_type)
-
 def log_db_connection_error(connection_type: str, error_message: str) -> None:
     log(
         "mito_ai_db_connection_error", 
@@ -354,7 +345,33 @@ def log_db_connection_error(connection_type: str, error_message: str) -> None:
             "error_message": error_message,
         }
     )
+
+def log_file_upload_attempt(
+    filename: str, file_extension: str, is_chunked: bool, total_chunks: int
+) -> None:
+    log(
+        "mito_ai_file_upload_attempt",
+        params={
+            "filename": filename,
+            "file_extension": file_extension,
+            "is_chunked": is_chunked,
+            "total_chunks": total_chunks,
+        },
+    )
+
+def log_file_upload_failure(error: str) -> None:
+    log("mito_ai_file_upload_failure", params={"error_message": error})
+
+def log_ai_completion_retry(key_type: Literal['mito_server_key', 'user_key'], message_type: MessageType, error: BaseException) -> None:
+    log(MITO_AI_COMPLETION_RETRY, params={KEY_TYPE_PARAM: key_type, "message_type": message_type}, key_type=key_type, error=error)
     
+def log_ai_completion_error(key_type: Literal['mito_server_key', 'user_key'], message_type: MessageType, error: BaseException) -> None:
+    log(MITO_AI_COMPLETION_ERROR, params={KEY_TYPE_PARAM: key_type, "message_type": message_type}, key_type=key_type, error=error)
+    
+def log_mito_server_free_tier_limit_reached(key_type: Literal['mito_server_key', 'user_key'], message_type: MessageType) -> None:
+    log(MITO_SERVER_FREE_TIER_LIMIT_REACHED, params={KEY_TYPE_PARAM: key_type, "message_type": message_type}, key_type=key_type)
+
+
 #################################
 # Streamlit Conversion
 #################################