mito-ai 0.1.50__py3-none-any.whl

mito_ai/tests/message_history/test_message_history_utils.py

@@ -0,0 +1,469 @@
+# Copyright (c) Saga Inc.
+# Distributed under the terms of the GNU Affero General Public License v3.0 License.
+
+import pytest
+from typing import Callable, List, cast
+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 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 (
+    AgentExecutionMetadata,
+    AgentSmartDebugMetadata,
+    AIOptimizedCell,
+    ThreadID,
+)
+from mito_ai.completions.prompt_builders.prompt_constants import (
+    FILES_SECTION_HEADING,
+    VARIABLES_SECTION_HEADING,
+    CODE_SECTION_HEADING,
+    ACTIVE_CELL_ID_SECTION_HEADING,
+    JUPYTER_NOTEBOOK_SECTION_HEADING,
+    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"]
+TEST_CODE = "import pandas as pd\ndf = pd.read_csv('data.csv')"
+TEST_INPUT = "Calculate the mean of col1"
+TEST_ERROR = "AttributeError: 'Series' object has no attribute 'mena'"
+
+def test_trim_sections_basic() -> None:
+    """Test trimming sections on a simple string with all section types."""
+    content = f"""Some text before.
+
+{FILES_SECTION_HEADING}
+file1.csv
+file2.txt
+file3.py
+
+{VARIABLES_SECTION_HEADING}
+var1 = 1
+var2 = "string"
+var3 = [1, 2, 3]
+
+{JUPYTER_NOTEBOOK_SECTION_HEADING}
+[
+    {{
+        "cell_type": "code",
+        "id": "cell1",
+        "code": "print('hello world')"
+    }}
+]
+
+{ACTIVE_CELL_ID_SECTION_HEADING}
+cell1
+
+{CODE_SECTION_HEADING}
+```python
+def hello():
+    print("world")
+```
+
+Some text after."""
+
+    result = trim_sections_from_message_content(content)
+
+    # Verify sections are replaced with placeholders
+    assert f"{FILES_SECTION_HEADING} {CONTENT_REMOVED_PLACEHOLDER}" in result
+    assert f"{VARIABLES_SECTION_HEADING} {CONTENT_REMOVED_PLACEHOLDER}" in result
+    assert f"{JUPYTER_NOTEBOOK_SECTION_HEADING} {CONTENT_REMOVED_PLACEHOLDER}" in result
+    assert f"{ACTIVE_CELL_ID_SECTION_HEADING} {CONTENT_REMOVED_PLACEHOLDER}" in result
+    
+    # Verify sections are not in the result anymore
+    assert "file1.csv" not in result
+    assert "var1 = 1" not in result
+    assert "cell_type" not in result
+    assert "cell1" not in result
+    
+    # Verify other content is preserved
+    assert "Some text before." in result
+    assert "Some text after." in result
+    assert f"{CODE_SECTION_HEADING}" in result
+    assert "def hello():" in result
+
+
+# Parameterized test cases for prompt builders (except inline completer)
+PROMPT_BUILDER_TEST_CASES = [
+    # Chat prompt
+    (
+        lambda: create_chat_prompt(TEST_VARIABLES, TEST_FILES, TEST_CODE, "cell1", False, TEST_INPUT),
+        ["Your task: Calculate the mean of col1"],
+        ["data.csv\nscript.py", f"{VARIABLES_SECTION_HEADING}\n'df': pd.DataFrame"],
+    ),
+    # Agent execution prompt
+    (
+        lambda: create_agent_execution_prompt(
+            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
+            )
+        ),
+        ["Your task: \nCalculate the mean of col1"],
+        ["data.csv\nscript.py", f"'df': pd.DataFrame", "import pandas as pd"],
+    ),
+    # Smart debug prompt
+    (
+        lambda: create_error_prompt(TEST_ERROR, TEST_CODE, "cell1", TEST_VARIABLES, TEST_FILES),
+        ["Error Traceback:", TEST_ERROR],
+        [
+            f"{FILES_SECTION_HEADING}\ndata.csv\nscript.py",
+            f"{VARIABLES_SECTION_HEADING}\n'df': pd.DataFrame",
+            f"{ACTIVE_CELL_ID_SECTION_HEADING}\ncell1",
+        ],
+    ),
+    # Explain code prompt (doesn't have sections to trim)
+    (
+        lambda: create_explain_code_prompt(TEST_CODE),
+        ["import pandas as pd"],
+        [],
+    ),
+    # Agent smart debug prompt
+    (
+        lambda: create_agent_smart_debug_prompt(
+            AgentSmartDebugMetadata(
+                variables=TEST_VARIABLES,
+                files=TEST_FILES,
+                aiOptimizedCells=[
+                    AIOptimizedCell(cell_type="code", id="cell1", code=TEST_CODE)
+                ],
+                error_message_producing_code_cell_id="cell1",
+                errorMessage=TEST_ERROR,
+                promptType="agent:autoErrorFixup",
+                threadId=ThreadID("test-thread-id"),
+                isChromeBrowser=True
+            )
+        ),
+        ["Error Traceback:", TEST_ERROR],
+        [
+            f"{FILES_SECTION_HEADING}\n{TEST_FILES[0]}",
+            f"{VARIABLES_SECTION_HEADING}\n{TEST_VARIABLES[0]}",
+            "cell_type",
+        ],
+    ),
+]
+
+
+@pytest.mark.parametrize("prompt_builder,expected_in_result,expected_not_in_result", PROMPT_BUILDER_TEST_CASES)
+def test_prompt_builder_trimming(prompt_builder: Callable[[], str], expected_in_result: List[str], expected_not_in_result: List[str]) -> None:
+    """Test trimming for different prompt builders."""
+    # Create prompt using the provided builder function
+    content = prompt_builder()
+
+    # Trim the content
+    result = trim_sections_from_message_content(content)
+
+    # If none of the section headings are present, the content shouldn't change
+    has_sections_to_trim = any(
+        heading in content 
+        for heading in [FILES_SECTION_HEADING, VARIABLES_SECTION_HEADING, JUPYTER_NOTEBOOK_SECTION_HEADING, ACTIVE_CELL_ID_SECTION_HEADING]
+    )
+    
+    if not has_sections_to_trim:
+        assert result == content
+        # Verify expected content is still in the result
+        for expected in expected_in_result:
+            assert expected in result
+        return
+
+    # Check for each section if it was present in the original content
+    if FILES_SECTION_HEADING in content:
+        assert f"{FILES_SECTION_HEADING} {CONTENT_REMOVED_PLACEHOLDER}" in result
+        
+    if VARIABLES_SECTION_HEADING in content:
+        assert f"{VARIABLES_SECTION_HEADING} {CONTENT_REMOVED_PLACEHOLDER}" in result
+        
+    if JUPYTER_NOTEBOOK_SECTION_HEADING in content:
+        assert (
+            f"{JUPYTER_NOTEBOOK_SECTION_HEADING} {CONTENT_REMOVED_PLACEHOLDER}"
+            in result
+        )
+
+    if ACTIVE_CELL_ID_SECTION_HEADING in content:
+        assert f"{ACTIVE_CELL_ID_SECTION_HEADING} {CONTENT_REMOVED_PLACEHOLDER}" in result
+
+    # Verify expected content is still in the result
+    for expected in expected_in_result:
+        assert expected in result
+
+    # Verify content that should be removed is not in the result
+    for not_expected in expected_not_in_result:
+        assert not_expected not in result
+
+def test_no_sections_to_trim() -> None:
+    """Test trimming content with no sections to trim."""
+    content = "This is a simple message with no sections to trim."
+    result = trim_sections_from_message_content(content)
+    assert result == content
+
+
+# Tests for trim_old_messages function
+def test_trim_old_messages_only_trims_user_messages() -> None:
+    """Test that trim_old_messages only trims content from user messages."""
+    # Create test messages with different roles
+    user_message_with_sections = f"""User prompt with sections.
+{FILES_SECTION_HEADING}
+file1.csv
+file2.txt
+"""
+    system_message_with_sections = f"""System message with sections.
+{FILES_SECTION_HEADING}
+file1.csv
+file2.txt
+"""
+    assistant_message_with_sections = f"""Assistant message with sections.
+{FILES_SECTION_HEADING}
+file1.csv
+file2.txt
+"""
+    
+    # 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 1"},
+        {"role": "user", "content": "Recent user message 2"},
+        {"role": "user", "content": "Recent user message 3"},
+    ]
+    
+    result = trim_old_messages(messages)
+    
+    # System message should remain unchanged even though it's old
+    system_content = result[0].get("content")
+    assert isinstance(system_content, str)
+    assert system_content == system_message_with_sections
+    assert FILES_SECTION_HEADING in system_content
+    assert "file1.csv" in system_content
+    
+    # First user message should be trimmed
+    assert result[1]["role"] == "user"
+    user_content = result[1].get("content")
+    assert isinstance(user_content, str)
+    assert f"{FILES_SECTION_HEADING} {CONTENT_REMOVED_PLACEHOLDER}" in user_content
+    assert "file1.csv" not in user_content
+    
+    # Assistant message should remain unchanged even though it's old
+    assistant_content = result[2].get("content")
+    assert isinstance(assistant_content, str)
+    assert assistant_content == assistant_message_with_sections
+    assert FILES_SECTION_HEADING in assistant_content
+    assert "file1.csv" in assistant_content
+    
+    # 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 MESSAGE_HISTORY_TRIM_THRESHOLD."""
+    # Create test messages
+    old_message_1 = f"""Old message 1.
+{FILES_SECTION_HEADING}
+file1.csv
+"""
+    old_message_2 = f"""Old message 2.
+{FILES_SECTION_HEADING}
+file2.csv
+"""
+    recent_message_1 = f"""Recent message 1.
+{FILES_SECTION_HEADING}
+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
+    messages: List[ChatCompletionMessageParam] = [
+        {"role": "user", "content": old_message_1},
+        {"role": "user", "content": old_message_2},
+        {"role": "user", "content": recent_message_1},
+        {"role": "user", "content": recent_message_2},
+        {"role": "user", "content": recent_message_3},
+    ]
+    
+    # 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")
+    assert isinstance(old_content_1, str)
+    assert f"{FILES_SECTION_HEADING} {CONTENT_REMOVED_PLACEHOLDER}" in old_content_1
+    assert "file1.csv" not in old_content_1
+    
+    old_content_2 = result[1].get("content")
+    assert isinstance(old_content_2, str)
+    assert f"{FILES_SECTION_HEADING} {CONTENT_REMOVED_PLACEHOLDER}" in old_content_2
+    assert "file2.csv" not in old_content_2
+    
+    # Recent messages should remain unchanged
+    recent_content_1 = result[2].get("content")
+    assert isinstance(recent_content_1, str)
+    assert recent_content_1 == recent_message_1
+    assert FILES_SECTION_HEADING in recent_content_1
+    assert "file3.csv" in recent_content_1
+    
+    recent_content_2 = result[3].get("content")
+    assert isinstance(recent_content_2, str)
+    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)
+    assert result == []
+
+
+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)
+    
+    # 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"
+    
+    assistant_content = result[1].get("content")
+    assert isinstance(assistant_content, str)
+    assert assistant_content == "Assistant message 1"
+
+
+def test_trim_mixed_content_messages() -> None:
+    """
+    Tests that when a message contains sections other than text (like image_url),
+    those sections are removed completely, leaving only the text content.
+    """
+    # Create sample message with mixed content (text and image)
+    mixed_content_message = cast(ChatCompletionMessageParam, {
+        "role": "user",
+        "content": [
+            {
+                "type": "text",
+                "text": "What is in this image?"
+            },
+            {
+                "type": "image_url",
+                "image_url": {"url": "data:image/png;base64,someimagedata"}
+            }
+        ]
+    })
+    
+    # Create sample message list with one old message (the mixed content)
+    # 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": "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)
+    
+    # Verify that the first message has been trimmed properly
+    assert trimmed_messages[0]["role"] == "user"
+    assert trimmed_messages[0]["content"] == "What is in this image?"
+    
+    # Verify that the recent messages are untouched
+    assert trimmed_messages[1] == message_list[1]
+    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

@@ -0,0 +1,152 @@
+# Copyright (c) Saga Inc.
+# Distributed under the terms of the GNU Affero General Public License v3.0 License.
+
+import pytest
+from datetime import datetime
+from unittest.mock import patch
+from mito_ai.utils.server_limits import (
+    check_mito_server_quota,
+    OS_MONTHLY_AI_COMPLETIONS_LIMIT,
+    OS_MONTHLY_AUTOCOMPLETE_LIMIT,
+)
+from mito_ai.completions.models import MessageType
+from mito_ai.utils.open_ai_utils import _prepare_request_data_and_headers
+
+REALLY_OLD_DATE = "2020-01-01"
+TODAY = datetime.now().strftime("%Y-%m-%d")
+
+
+def test_check_mito_server_quota_open_source_user() -> None:
+    # Under chat completions limit
+    with patch("mito_ai.utils.server_limits.get_chat_completion_count", return_value=1) as mock_count, \
+         patch("mito_ai.utils.server_limits.get_last_reset_date", return_value=TODAY) as mock_date, \
+         patch("mito_ai.utils.server_limits.is_pro", return_value=False):
+        
+        check_mito_server_quota(MessageType.CHAT)
+        assert mock_count.called
+        assert mock_date.called
+        assert mock_count.return_value == 1
+        assert mock_date.return_value == TODAY
+
+    # Under autocomplete limit
+    with patch("mito_ai.utils.server_limits.get_autocomplete_count", return_value=1) as mock_count, \
+         patch("mito_ai.utils.server_limits.get_last_reset_date", return_value=TODAY) as mock_date, \
+         patch("mito_ai.utils.server_limits.is_pro", return_value=False):
+        
+        check_mito_server_quota(MessageType.INLINE_COMPLETION)
+        assert mock_count.called
+        assert mock_date.called
+        assert mock_count.return_value == 1
+        assert mock_date.return_value == TODAY
+
+    # Over chat completions limit
+    with pytest.raises(PermissionError), \
+         patch("mito_ai.utils.server_limits.get_chat_completion_count", return_value=OS_MONTHLY_AI_COMPLETIONS_LIMIT + 1) as mock_count, \
+         patch("mito_ai.utils.server_limits.get_last_reset_date", return_value=TODAY) as mock_date, \
+         patch("mito_ai.utils.server_limits.is_pro", return_value=False):
+        
+        check_mito_server_quota(MessageType.CHAT)
+        assert mock_count.called
+        assert mock_date.called
+        assert mock_count.return_value == OS_MONTHLY_AI_COMPLETIONS_LIMIT + 1
+        assert mock_date.return_value == TODAY
+
+    # Over autocomplete limit
+    with pytest.raises(PermissionError), \
+         patch("mito_ai.utils.server_limits.get_autocomplete_count", return_value=OS_MONTHLY_AUTOCOMPLETE_LIMIT + 1) as mock_count, \
+         patch("mito_ai.utils.server_limits.get_last_reset_date", return_value=TODAY) as mock_date, \
+         patch("mito_ai.utils.server_limits.is_pro", return_value=False):
+        
+        check_mito_server_quota(MessageType.INLINE_COMPLETION)
+        assert mock_count.called
+        assert mock_date.called
+        assert mock_count.return_value == OS_MONTHLY_AUTOCOMPLETE_LIMIT + 1
+        assert mock_date.return_value == TODAY
+
+
+def test_check_mito_server_quota_pro_user() -> None:
+    # No error should be thrown since pro users don't have limits
+    with patch("mito_ai.utils.server_limits.is_pro", return_value=True), \
+         patch("mito_ai.utils.server_limits.get_chat_completion_count", return_value=1000), \
+         patch("mito_ai.utils.server_limits.get_last_reset_date", return_value=REALLY_OLD_DATE):
+        
+        check_mito_server_quota(MessageType.CHAT)
+
+def test_prepare_request_data_and_headers_basic() -> None:
+    """Test basic functionality of _prepare_request_data_and_headers"""
+
+    # Mock the user fields
+    with patch("mito_ai.utils.open_ai_utils.get_user_field") as mock_get_user_field:
+        mock_get_user_field.side_effect = ["test@example.com", "user123"]
+        
+        # Mock the quota check
+        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"""
+    with patch("mito_ai.utils.open_ai_utils.get_user_field") as mock_get_user_field:
+        mock_get_user_field.side_effect = ["test@example.com", "user123"]
+        
+        with patch("mito_ai.utils.open_ai_utils.check_mito_server_quota"):
+            data, _ = _prepare_request_data_and_headers(
+                last_message_content=None,
+                ai_completion_data={},
+                timeout=30,
+                max_retries=3,
+                message_type=MessageType.CHAT
+            )
+
+            # Verify empty string is used for null message
+            assert data["user_input"] == ""
+
+def test_prepare_request_data_and_headers_caches_user_info() -> None:
+    """Test that user info is cached after first call"""
+    # Mock both the global variables and the get_user_field function
+    with patch("mito_ai.utils.open_ai_utils.__user_email", None), \
+         patch("mito_ai.utils.open_ai_utils.__user_id", None), \
+         patch("mito_ai.utils.open_ai_utils.get_user_field") as mock_get_user_field:
+        
+        mock_get_user_field.side_effect = ["test@example.com", "user123"]
+        
+        with patch("mito_ai.utils.open_ai_utils.check_mito_server_quota"):
+            # First call
+            data1, _ = _prepare_request_data_and_headers(
+                last_message_content="test",
+                ai_completion_data={},
+                timeout=30,
+                max_retries=3,
+                message_type=MessageType.CHAT
+            )
+
+            # Second call
+            data2, _ = _prepare_request_data_and_headers(
+                last_message_content="test",
+                ai_completion_data={},
+                timeout=30,
+                max_retries=3,
+                message_type=MessageType.CHAT
+            )
+
+            # Verify get_user_field was only called twice (once for email, once for user_id)
+            assert mock_get_user_field.call_count == 2
+            
+            # Verify both calls return same user info
+            assert data1["email"] == data2["email"] == "test@example.com"
+            assert data1["user_id"] == data2["user_id"] == "user123"