mito-ai 0.1.54__py3-none-any.whl → 0.1.56__py3-none-any.whl

mito_ai/tests/message_history/test_message_history_utils.py

@@ -2,391 +2,311 @@
 # Distributed under the terms of the GNU Affero General Public License v3.0 License.
 
 import pytest
-from typing import Callable, List, cast
+from typing import 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 mito_ai.utils.message_history_utils import trim_message_content, trim_old_messages
 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")
-```
+from mito_ai.completions.models import ThreadID
 
-Some text after."""
 
-    result = trim_sections_from_message_content(content)
+# Tests for trim_message_content function
+
+def test_trim_message_content_removes_sections_based_on_threshold() -> None:
+    """Test that sections are removed when message_age >= trim_after_messages threshold."""
+    # FilesSection has trim_after_messages = 3
+    # VariablesSection has trim_after_messages = 6
+    # NotebookSection has trim_after_messages = 6
+    
+    content = """Some text before.
+
+<Files>file1.csv
+file2.txt</Files>
+
+<Variables>var1 = 1
+var2 = "string"</Variables>
+
+<Notebook>[
+    {{"cell_type": "code", "id": "cell1"}}
+]</Notebook>
+
+Some text after."""
 
-    # 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
+    # Test with message_age = 2 (should NOT trim Files, Variables, or Notebook)
+    result = trim_message_content(content, message_age=2)
+    assert "<Files>" in result
+    assert "file1.csv" in result
+    assert "<Variables>" in result
+    assert "var1 = 1" in result
+    assert "<Notebook>" in result
+    assert "cell_type" in result
+    
+    # Test with message_age = 3 (should trim Files, but NOT Variables or Notebook)
+    result = trim_message_content(content, message_age=3)
+    assert "<Files>" not in result
+    assert "file1.csv" not in result
+    assert "<Variables>" in result
+    assert "var1 = 1" in result
+    assert "<Notebook>" in result
+    assert "cell_type" in result
     
-    # Verify sections are not in the result anymore
+    # Test with message_age = 6 (should trim Files, Variables, and Notebook)
+    result = trim_message_content(content, message_age=6)
+    assert "<Files>" not in result
     assert "file1.csv" not in result
+    assert "<Variables>" not in result
     assert "var1 = 1" not in result
+    assert "<Notebook>" 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."""
+
+
+def test_trim_message_content_removes_nested_xml_tags() -> None:
+    """Test that trimming correctly removes nested XML tags within a section."""
+    # ExampleSection can contain nested XML tags like <Files> and <Variables>
+    # ExampleSection has trim_after_messages = 3
+    
+    content = """Some text before.
+
+<Example name="Example 1">
+<Files>file1.csv
+file2.txt</Files>
+<Variables>var1 = 1</Variables>
+Some example text here.
+</Example>
+
+Some text after."""
+
+    # Test with message_age = 2 (should NOT trim Example)
+    result = trim_message_content(content, message_age=2)
+    assert "<Example" in result
+    assert "file1.csv" in result
+    assert "<Variables>" in result
+    assert "var1 = 1" in result
+    
+    # Test with message_age = 3 (should trim entire Example including nested tags)
+    result = trim_message_content(content, message_age=3)
+    assert "<Example" not in result
+    assert "file1.csv" not in result
+    assert "<Variables>" not in result
+    assert "var1 = 1" not in result
+    assert "Some example text here." not in result
+    
+    # Verify other content is preserved
+    assert "Some text before." in result
+    assert "Some text after." in result
+
+
+def test_trim_message_content_handles_multiple_sections_of_same_type() -> None:
+    """Test that all instances of a section type are removed when threshold is met."""
+    content = """Text before.
+
+<Files>file1.csv</Files>
+
+Some middle text.
+
+<Files>file2.txt
+file3.py</Files>
+
+Text after."""
+
+    # Test with message_age = 3 (should remove all Files sections)
+    result = trim_message_content(content, message_age=3)
+    assert "<Files>" not in result
+    assert "file1.csv" not in result
+    assert "file2.txt" not in result
+    assert "file3.py" not in result
+    
+    # Verify other content is preserved
+    assert "Text before." in result
+    assert "Some middle text." in result
+    assert "Text after." in result
+
+
+def test_trim_message_content_preserves_sections_with_none_threshold() -> None:
+    """Test that sections with trim_after_messages = None are never trimmed."""
+    # ActiveCellIdSection has trim_after_messages = None
+    # TaskSection has trim_after_messages = None
+    
+    content = """Some text.
+
+<ActiveCellId>cell1</ActiveCellId>
+
+<Task>Your task: Do something</Task>
+
+<Files>file1.csv</Files>
+
+More text."""
+
+    # Test with very high message_age
+    result = trim_message_content(content, message_age=100)
+    
+    # ActiveCellId and Task should remain (threshold = None)
+    assert "<ActiveCellId>" in result
+    assert "cell1" in result
+    assert "<Task>" in result
+    assert "Your task: Do something" in result
+    
+    # Files should be removed (threshold = 3)
+    assert "<Files>" not in result
+    assert "file1.csv" not in result
+
+
+def test_trim_message_content_handles_empty_content() -> None:
+    """Test that trimming handles empty content correctly."""
+    content = ""
+    result = trim_message_content(content, message_age=5)
+    assert result == ""
+
+
+def test_trim_message_content_handles_content_without_sections() -> None:
+    """Test that trimming preserves content without XML sections."""
     content = "This is a simple message with no sections to trim."
-    result = trim_sections_from_message_content(content)
+    result = trim_message_content(content, message_age=5)
     assert result == content
 
 
+def test_trim_message_content_handles_malformed_xml_gracefully() -> None:
+    """Test that trimming handles malformed XML gracefully."""
+    # Content with unclosed tags or mismatched tags
+    content = """Some text.
+
+<Files>file1.csv
+<Variables>var1 = 1</Variables>
+
+More text."""
+
+    # Should not crash, and should attempt to trim what it can
+    result = trim_message_content(content, message_age=3)
+    # The behavior depends on regex matching, but should not raise an exception
+    assert isinstance(result, str)
+
+
 # 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
-"""
+
+def test_trim_old_messages_calculates_message_age_correctly() -> None:
+    """Test that message age is calculated correctly based on position in message list."""
+    # Message age = total_messages - i - 1
+    # So for a list of 5 messages (indices 0-4):
+    # - Index 0 (oldest): age = 5 - 0 - 1 = 4
+    # - Index 1: age = 5 - 1 - 1 = 3
+    # - Index 2: age = 5 - 2 - 1 = 2
+    # - Index 3: age = 5 - 3 - 1 = 1
+    # - Index 4 (newest): age = 5 - 4 - 1 = 0
+    
+    # FilesSection has trim_after_messages = 3, so it should be trimmed when age >= 3
+    # VariablesSection has trim_after_messages = 6, so it should be trimmed when age >= 6
     
-    # 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"},
+        {"role": "user", "content": "Message 0 with <Files>file0.csv</Files> and <Variables>var0</Variables>"},
+        {"role": "user", "content": "Message 1 with <Files>file1.csv</Files> and <Variables>var1</Variables>"},
+        {"role": "user", "content": "Message 2 with <Files>file2.csv</Files> and <Variables>var2</Variables>"},
+        {"role": "user", "content": "Message 3 with <Files>file3.csv</Files> and <Variables>var3</Variables>"},
+        {"role": "user", "content": "Message 4 with <Files>file4.csv</Files> and <Variables>var4</Variables>"},
     ]
     
     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
+    # Message 0 (age=4): Files should be trimmed (4 >= 3), Variables should remain (4 < 6)
+    content_0 = result[0].get("content")
+    assert isinstance(content_0, str)
+    assert "<Files>" not in content_0
+    assert "file0.csv" not in content_0
+    assert "<Variables>" in content_0
+    assert "var0" in content_0
     
-    # 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
+    # Message 1 (age=3): Files should be trimmed (3 >= 3), Variables should remain (3 < 6)
+    content_1 = result[1].get("content")
+    assert isinstance(content_1, str)
+    assert "<Files>" not in content_1
+    assert "file1.csv" not in content_1
+    assert "<Variables>" in content_1
+    assert "var1" in content_1
     
-    # 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"
+    # Message 2 (age=2): Files should remain (2 < 3), Variables should remain (2 < 6)
+    content_2 = result[2].get("content")
+    assert isinstance(content_2, str)
+    assert "<Files>" in content_2
+    assert "file2.csv" in content_2
+    assert "<Variables>" in content_2
+    assert "var2" in content_2
     
-    recent_content_2 = result[4].get("content")
-    assert isinstance(recent_content_2, str)
-    assert recent_content_2 == "Recent user message 2"
+    # Message 3 (age=1): Files should remain (1 < 3), Variables should remain (1 < 6)
+    content_3 = result[3].get("content")
+    assert isinstance(content_3, str)
+    assert "<Files>" in content_3
+    assert "file3.csv" in content_3
+    assert "<Variables>" in content_3
+    assert "var3" in content_3
     
-    recent_content_3 = result[5].get("content")
-    assert isinstance(recent_content_3, str)
-    assert recent_content_3 == "Recent user message 3"
+    # Message 4 (age=0, newest): Files should remain (0 < 3), Variables should remain (0 < 6)
+    content_4 = result[4].get("content")
+    assert isinstance(content_4, str)
+    assert "<Files>" in content_4
+    assert "file4.csv" in content_4
+    assert "<Variables>" in content_4
+    assert "var4" in content_4
 
 
-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
+def test_trim_old_messages_only_trims_user_messages() -> None:
+    """Test that trim_old_messages only trims content from user messages."""
+    user_message_with_sections = """User prompt with sections.
+<Files>file1.csv
+file2.txt</Files>
 """
-    recent_message_2 = f"""Recent message 2.
-{FILES_SECTION_HEADING}
-file4.csv
+    system_message_with_sections = """System message with sections.
+<Files>file1.csv
+file2.txt</Files>
 """
-    recent_message_3 = f"""Recent message 3.
-{FILES_SECTION_HEADING}
-file5.csv
+    assistant_message_with_sections = """Assistant message with sections.
+<Files>file1.csv
+file2.txt</Files>
 """
     
-    # 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},
+        {"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"},
     ]
     
-    # 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)
+    # System message should remain unchanged (not trimmed)
+    system_content = result[0].get("content")
+    assert isinstance(system_content, str)
+    assert system_content == system_message_with_sections
+    assert "<Files>" in system_content
+    assert "file1.csv" in system_content
     
-    # Messages should remain unchanged since we have fewer than MESSAGE_HISTORY_TRIM_THRESHOLD (3) messages
-    user_content = result[0].get("content")
+    # User message should be trimmed (age=1, but Files threshold is 3, so actually not trimmed in this case)
+    # But let's test with an older message to ensure trimming works
+    user_content = result[1].get("content")
     assert isinstance(user_content, str)
-    assert user_content == "User message 1"
+    # With age=2, Files should remain (2 < 3)
+    assert "<Files>" in user_content
     
-    assistant_content = result[1].get("content")
+    # Assistant message should remain unchanged (not trimmed)
+    assistant_content = result[2].get("content")
     assert isinstance(assistant_content, str)
-    assert assistant_content == "Assistant message 1"
+    assert assistant_content == assistant_message_with_sections
+    assert "<Files>" in assistant_content
 
 
-def test_trim_mixed_content_messages() -> None:
+def test_trim_old_messages_handles_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.
+    those sections are removed completely, leaving only the trimmed 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?"
+                "text": "What is in this image? <Files>file1.csv</Files>"
             },
             {
                 "type": "image_url",
@@ -395,30 +315,43 @@ def test_trim_mixed_content_messages() -> None:
         ]
     })
     
-    # Create sample message list with one old message (the mixed content)
-    # and enough recent messages to exceed MESSAGE_HISTORY_TRIM_THRESHOLD (3)
+    # Create message list with old message (should be trimmed) and recent messages
     message_list: List[ChatCompletionMessageParam] = [
-        mixed_content_message,  # This should get trimmed
+        mixed_content_message,  # Age = 4, Files should be trimmed (4 >= 3)
         {"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
+        {"role": "user", "content": "Can you explain more?"},
+        {"role": "user", "content": "Another recent message"},
+        {"role": "user", "content": "Yet another recent message"}
     ]
     
-    # 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?"
+    first_content = trimmed_messages[0].get("content")
+    assert isinstance(first_content, list)
+    # Find the text section
+    text_section = next((s for s in first_content if s.get("type") == "text"), None)
+    assert text_section is not None
+    assert isinstance(text_section.get("text"), str)
+    assert "<Files>" not in text_section["text"]
+    assert "file1.csv" not in text_section["text"]
+    assert "What is in this image?" in text_section["text"]
     
-    # Verify that the recent messages are untouched
+    # Verify that 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_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_get_display_history_calls_update_last_interaction() -> None:
     """Test that get_display_history calls _update_last_interaction when retrieving a thread."""