mito-ai 0.1.33__py3-none-any.whl → 0.1.49__py3-none-any.whl

mito_ai/tests/deploy_app/test_app_deploy_utils.py

@@ -0,0 +1,89 @@
+# Copyright (c) Saga Inc.
+# Distributed under the terms of the GNU Affero General Public License v3.0 License.
+
+import zipfile
+import logging
+from mito_ai.app_deploy.app_deploy_utils import add_files_to_zip
+from mito_ai.path_utils import AbsoluteNotebookDirPath
+
+class TestAddFilesToZip:
+    """Test cases for add_files_to_zip helper function"""
+
+    def test_files_added_correctly(self, tmp_path):
+        """Ensure individual files are added correctly to the zip"""
+        # Create files
+        f1 = tmp_path / "file1.txt"
+        f1.write_text("file1 content")
+        f2 = tmp_path / "file2.txt"
+        f2.write_text("file2 content")
+
+        zip_path = tmp_path / "test.zip"
+        add_files_to_zip(str(zip_path), AbsoluteNotebookDirPath(str(tmp_path)), ["file1.txt", "file2.txt"], 'test-app-file-name.py')
+
+        with zipfile.ZipFile(zip_path, "r") as zf:
+            names = zf.namelist()
+            assert "file1.txt" in names
+            assert "file2.txt" in names
+            assert len(names) == 2
+            
+    def test_renames_app_file(self, tmp_path):
+        """Ensure individual files are added correctly to the zip"""
+        # Create files
+        f1 = tmp_path / "original-file-name.py"
+        f1.write_text("file1 content")
+        f2 = tmp_path / "file2.txt"
+        f2.write_text("file2 content")
+
+        zip_path = tmp_path / "test.zip"
+        add_files_to_zip(str(zip_path), AbsoluteNotebookDirPath(str(tmp_path)), ["original-file-name.py", "file2.txt"], 'original-file-name.py')
+
+        with zipfile.ZipFile(zip_path, "r") as zf:
+            names = zf.namelist()
+            assert "app.py" in names
+            assert "file2.txt" in names
+            assert len(names) == 2
+
+    def test_directories_added_recursively(self, tmp_path):
+        """Ensure directories are added recursively with correct relative paths"""
+        nested = tmp_path / "folder"
+        nested.mkdir()
+        (nested / "nested1.txt").write_text("nested1 content")
+        subfolder = nested / "sub"
+        subfolder.mkdir()
+        (subfolder / "nested2.txt").write_text("nested2 content")
+
+        zip_path = tmp_path / "test.zip"
+        add_files_to_zip(str(zip_path), AbsoluteNotebookDirPath(str(tmp_path)), ["folder"], 'test-app.py')
+
+        with zipfile.ZipFile(zip_path, "r") as zf:
+            names = zf.namelist()
+            assert "folder/nested1.txt" in names
+            assert "folder/sub/nested2.txt" in names
+
+    def test_missing_files_skipped(self, tmp_path, caplog):
+        """Ensure missing files do not break the function and warning is logged"""
+        caplog.set_level(logging.WARNING)
+        zip_path = tmp_path / "test.zip"
+        add_files_to_zip(str(zip_path), AbsoluteNotebookDirPath(str(tmp_path)), ["does_not_exist.txt"], 'test-app.py', logger=logging.getLogger())
+
+        # Zip should exist but be empty
+        with zipfile.ZipFile(zip_path, "r") as zf:
+            assert zf.namelist() == []
+
+        # Check warning was logged
+        assert any("Skipping missing file" in record.message for record in caplog.records)
+
+    def test_arcname_paths_correct(self, tmp_path):
+        """Ensure arcname paths inside zip preserve relative paths to base_path"""
+        (tmp_path / "file.txt").write_text("content")
+        folder = tmp_path / "folder"
+        folder.mkdir()
+        (folder / "nested.txt").write_text("nested content")
+
+        zip_path = tmp_path / "test.zip"
+        add_files_to_zip(str(zip_path), AbsoluteNotebookDirPath(str(tmp_path)), ["file.txt", "folder"], 'test-app.py')
+
+        with zipfile.ZipFile(zip_path, "r") as zf:
+            names = zf.namelist()
+            assert "file.txt" in names
+            assert "folder/nested.txt" in names

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,282 @@
+# 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]
+
+
+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/tests/message_history/test_generate_short_chat_name.py

@@ -6,10 +6,6 @@ from unittest.mock import AsyncMock, MagicMock, patch
 from traitlets.config import Config
 from mito_ai.completions.message_history import generate_short_chat_name
 from mito_ai.completions.providers import OpenAIProvider
-from mito_ai.completions.models import MessageType
-from mito_ai.openai_client import OPENAI_FAST_MODEL
-from mito_ai.anthropic_client import ANTHROPIC_FAST_MODEL  
-from mito_ai.gemini_client import GEMINI_FAST_MODEL
 
 
 @pytest.fixture

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"]
@@ -103,12 +106,15 @@ PROMPT_BUILDER_TEST_CASES = [
             AgentExecutionMetadata(
                 variables=TEST_VARIABLES,
                 files=TEST_FILES,
+                notebookPath='/test-notebook-path.ipynb',
+                notebookID='test-notebook-id',
                 aiOptimizedCells=[
                     AIOptimizedCell(cell_type="code", id="cell1", code=TEST_CODE)
                 ],
                 input=TEST_INPUT,
                 promptType="agent:execution",
                 threadId=ThreadID("test-thread-id"),
+                activeCellId="cell1",
                 isChromeBrowser=True
             )
         ),
@@ -230,16 +236,17 @@ file1.csv
 file2.txt
 """
     
-    # Create test messages with proper typing
+    # Create test messages with proper typing 
     messages: List[ChatCompletionMessageParam] = [
         {"role": "system", "content": system_message_with_sections},
         {"role": "user", "content": user_message_with_sections},
         {"role": "assistant", "content": assistant_message_with_sections},
-        {"role": "user", "content": "Recent user message"},
+        {"role": "user", "content": "Recent user message 1"},
+        {"role": "user", "content": "Recent user message 2"},
+        {"role": "user", "content": "Recent user message 3"},
     ]
     
-    # Keep only the most recent message
-    result = trim_old_messages(messages, keep_recent=1)
+    result = trim_old_messages(messages)
     
     # System message should remain unchanged even though it's old
     system_content = result[0].get("content")
@@ -262,14 +269,22 @@ file2.txt
     assert FILES_SECTION_HEADING in assistant_content
     assert "file1.csv" in assistant_content
     
-    # Recent user message should remain unchanged
-    recent_content = result[3].get("content")
-    assert isinstance(recent_content, str)
-    assert recent_content == "Recent user message"
+    # Recent user messages should remain unchanged
+    recent_content_1 = result[3].get("content")
+    assert isinstance(recent_content_1, str)
+    assert recent_content_1 == "Recent user message 1"
+    
+    recent_content_2 = result[4].get("content")
+    assert isinstance(recent_content_2, str)
+    assert recent_content_2 == "Recent user message 2"
+    
+    recent_content_3 = result[5].get("content")
+    assert isinstance(recent_content_3, str)
+    assert recent_content_3 == "Recent user message 3"
 
 
 def test_trim_old_messages_preserves_recent_messages() -> None:
-    """Test that trim_old_messages preserves the most recent messages based on keep_recent parameter."""
+    """Test that trim_old_messages preserves the most recent messages based on MESSAGE_HISTORY_TRIM_THRESHOLD."""
     # Create test messages
     old_message_1 = f"""Old message 1.
 {FILES_SECTION_HEADING}
@@ -286,6 +301,10 @@ file3.csv
     recent_message_2 = f"""Recent message 2.
 {FILES_SECTION_HEADING}
 file4.csv
+"""
+    recent_message_3 = f"""Recent message 3.
+{FILES_SECTION_HEADING}
+file5.csv
 """
     
     # Create test messages with proper typing
@@ -294,10 +313,11 @@ file4.csv
         {"role": "user", "content": old_message_2},
         {"role": "user", "content": recent_message_1},
         {"role": "user", "content": recent_message_2},
+        {"role": "user", "content": recent_message_3},
     ]
     
-    # Keep the 2 most recent messages
-    result = trim_old_messages(messages, keep_recent=2)
+    # Test with MESSAGE_HISTORY_TRIM_THRESHOLD (3) - only the first 2 messages should be trimmed
+    result = trim_old_messages(messages)
     
     # Old messages should be trimmed
     old_content_1 = result[0].get("content")
@@ -322,24 +342,30 @@ file4.csv
     assert recent_content_2 == recent_message_2
     assert FILES_SECTION_HEADING in recent_content_2
     assert "file4.csv" in recent_content_2
+    
+    recent_content_3 = result[4].get("content")
+    assert isinstance(recent_content_3, str)
+    assert recent_content_3 == recent_message_3
+    assert FILES_SECTION_HEADING in recent_content_3
+    assert "file5.csv" in recent_content_3
 
 def test_trim_old_messages_empty_list() -> None:
     """Test that trim_old_messages handles empty message lists correctly."""
     messages: List[ChatCompletionMessageParam] = []
-    result = trim_old_messages(messages, keep_recent=2)
+    result = trim_old_messages(messages)
     assert result == []
 
 
-def test_trim_old_messages_fewer_than_keep_recent() -> None:
-    """Test that trim_old_messages doesn't modify messages if there are fewer than keep_recent."""
+def test_trim_old_messages_fewer_than_threshold() -> None:
+    """Test that trim_old_messages doesn't modify messages if there are fewer than MESSAGE_HISTORY_TRIM_THRESHOLD."""
     messages: List[ChatCompletionMessageParam] = [
         {"role": "user", "content": "User message 1"},
         {"role": "assistant", "content": "Assistant message 1"},
     ]
     
-    result = trim_old_messages(messages, keep_recent=3)
+    result = trim_old_messages(messages)
     
-    # Messages should remain unchanged
+    # Messages should remain unchanged since we have fewer than MESSAGE_HISTORY_TRIM_THRESHOLD (3) messages
     user_content = result[0].get("content")
     assert isinstance(user_content, str)
     assert user_content == "User message 1"
@@ -370,15 +396,17 @@ def test_trim_mixed_content_messages() -> None:
     })
     
     # Create sample message list with one old message (the mixed content)
-    # and one recent message (to not be trimmed)
+    # and enough recent messages to exceed MESSAGE_HISTORY_TRIM_THRESHOLD (3)
     message_list: List[ChatCompletionMessageParam] = [
         mixed_content_message,  # This should get trimmed
         {"role": "assistant", "content": "That's a chart showing data trends"},
-        {"role": "user", "content": "Can you explain more?"}  # Recent message, should not be trimmed
+        {"role": "user", "content": "Can you explain more?"},  # Recent message, should not be trimmed
+        {"role": "user", "content": "Another recent message"},  # Recent message, should not be trimmed
+        {"role": "user", "content": "Yet another recent message"}  # Recent message, should not be trimmed
     ]
     
     # Apply the trimming function
-    trimmed_messages = trim_old_messages(message_list, keep_recent=2)
+    trimmed_messages = trim_old_messages(message_list)
     
     # Verify that the first message has been trimmed properly
     assert trimmed_messages[0]["role"] == "user"
@@ -386,4 +414,56 @@ 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]
+    assert trimmed_messages[3] == message_list[3]
+    assert trimmed_messages[4] == message_list[4]
+
+
+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/tests/open_ai_utils_test.py

@@ -80,28 +80,24 @@ def test_prepare_request_data_and_headers_basic() -> None:
         mock_get_user_field.side_effect = ["test@example.com", "user123"]
         
         # Mock the quota check
-        with patch("mito_ai.utils.open_ai_utils.check_mito_server_quota") as mock_check_quota:
-            data, headers = _prepare_request_data_and_headers(
-                last_message_content="test message",
-                ai_completion_data={"key": "value"},
-                timeout=30,
-                max_retries=3,
-                message_type=MessageType.CHAT
-            )
-
-            # Verify quota check was called
-            mock_check_quota.assert_called_once_with(MessageType.CHAT)
-
-            # Verify data structure
-            assert data["timeout"] == 30
-            assert data["max_retries"] == 3
-            assert data["email"] == "test@example.com"
-            assert data["user_id"] == "user123"
-            assert data["data"] == {"key": "value"}
-            assert data["user_input"] == "test message"
-
-            # Verify headers
-            assert headers == {"Content-Type": "application/json"}
+        data, headers = _prepare_request_data_and_headers(
+            last_message_content="test message",
+            ai_completion_data={"key": "value"},
+            timeout=30,
+            max_retries=3,
+            message_type=MessageType.CHAT
+        )
+
+        # Verify data structure
+        assert data["timeout"] == 30
+        assert data["max_retries"] == 3
+        assert data["email"] == "test@example.com"
+        assert data["user_id"] == "user123"
+        assert data["data"] == {"key": "value"}
+        assert data["user_input"] == "test message"
+
+        # Verify headers
+        assert headers == {"Content-Type": "application/json"}
 
 def test_prepare_request_data_and_headers_null_message() -> None:
     """Test handling of null message content"""