mito-ai 0.1.45__py3-none-any.whl → 0.1.47__py3-none-any.whl

mito_ai/tests/providers/test_anthropic_client.py

@@ -2,14 +2,14 @@
 # Distributed under the terms of the GNU Affero General Public License v3.0 License.
 
 import pytest
-from mito_ai.anthropic_client import get_anthropic_system_prompt_and_messages, extract_and_parse_anthropic_json_response, AnthropicClient
-from mito_ai.utils.anthropic_utils import get_anthropic_completion_function_params, FAST_ANTHROPIC_MODEL
+from mito_ai.anthropic_client import get_anthropic_system_prompt_and_messages, get_anthropic_system_prompt_and_messages_with_caching, add_cache_control_to_message, extract_and_parse_anthropic_json_response, AnthropicClient
+from mito_ai.utils.anthropic_utils import FAST_ANTHROPIC_MODEL
 from anthropic.types import Message, TextBlock, ToolUseBlock, Usage, ToolUseBlock, Message, Usage, TextBlock
 from openai.types.chat import ChatCompletionMessageParam, ChatCompletionUserMessageParam, ChatCompletionAssistantMessageParam, ChatCompletionSystemMessageParam
-from mito_ai.completions.models import MessageType, ResponseFormatInfo, AgentResponse
-from unittest.mock import MagicMock, patch
+from mito_ai.completions.models import MessageType
+from unittest.mock import patch
 import anthropic
-from typing import List, Dict, Any, cast, Union
+from typing import List, Dict, cast
 
 
 # Dummy base64 image (1x1 PNG)
@@ -272,4 +272,176 @@ async def test_model_selection_based_on_message_type(message_type, expected_mode
         # Verify that create was called with the expected model
         mock_create.assert_called_once()
         call_args = mock_create.call_args
-        assert call_args[1]['model'] == expected_model 
+        assert call_args[1]['model'] == expected_model
+
+
+# Caching Tests
+
+@pytest.mark.parametrize("message,expected_role,expected_content_type,expected_content_length,expected_cache_control", [
+    # String content message
+    (
+        {"role": "user", "content": "Hello world"},
+        "user",
+        list,
+        1,
+        True
+    ),
+    # List content message
+    (
+        {
+            "role": "user", 
+            "content": [
+                {"type": "text", "text": "First part"},
+                {"type": "text", "text": "Second part"}
+            ]
+        },
+        "user",
+        list,
+        2,
+        True
+    ),
+    # Empty content message
+    (
+        {"role": "user", "content": []},
+        "user",
+        list,
+        0,
+        False
+    ),
+    # Assistant message with string content
+    (
+        {"role": "assistant", "content": "I can help you with that."},
+        "assistant",
+        list,
+        1,
+        True
+    ),
+])
+def test_add_cache_control_to_message(message, expected_role, expected_content_type, expected_content_length, expected_cache_control):
+    """Test adding cache control to different types of messages."""
+    result = add_cache_control_to_message(message)
+    
+    assert result["role"] == expected_role
+    assert isinstance(result["content"], expected_content_type)
+    assert len(result["content"]) == expected_content_length
+    
+    if expected_cache_control and expected_content_length > 0:
+        # Should have cache_control on the last content block
+        last_block = result["content"][-1]
+        assert last_block["cache_control"] == {"type": "ephemeral"}
+        
+        # If there are multiple blocks, earlier blocks should not have cache_control
+        if expected_content_length > 1:
+            for i in range(expected_content_length - 1):
+                assert "cache_control" not in result["content"][i]
+    elif expected_content_length == 0:
+        # Empty content should return unchanged
+        assert result == message
+
+
+@pytest.mark.parametrize("messages,expected_system_type,expected_system_content", [
+    # With system prompt
+    (
+        [
+            ChatCompletionSystemMessageParam(role="system", content="You are a helpful assistant."),
+            ChatCompletionUserMessageParam(role="user", content="Hello!")
+        ],
+        list,
+        "You are a helpful assistant.",
+    ),
+    # Without system prompt
+    (
+        [
+            ChatCompletionUserMessageParam(role="user", content="Hello!"),
+            ChatCompletionAssistantMessageParam(role="assistant", content="Hi there!")
+        ],
+        anthropic.Omit,
+        None,
+    ),
+    # Multiple system messages (should take last one)
+    (
+        [
+            ChatCompletionSystemMessageParam(role="system", content="First system message."),
+            ChatCompletionSystemMessageParam(role="system", content="Second system message."),
+            ChatCompletionUserMessageParam(role="user", content="Hello!"),
+            ChatCompletionUserMessageParam(role="user", content="Hello!"),
+            ChatCompletionUserMessageParam(role="user", content="Hello!")
+        ],
+        list,
+        "Second system message.",
+    ),
+])
+def test_caching_system_prompt_scenarios(messages, expected_system_type, expected_system_content):
+    """Test caching with different system prompt scenarios."""
+    system_prompt, anthropic_messages = get_anthropic_system_prompt_and_messages_with_caching(messages)
+    
+    # Check system prompt
+    assert isinstance(system_prompt, expected_system_type)
+    if expected_system_content:
+        assert system_prompt[0]["text"] == expected_system_content
+        assert system_prompt[0]["cache_control"] == {"type": "ephemeral"}
+
+
+@pytest.mark.parametrize("message_count,expected_cache_boundary", [
+    (1, None),  # 1 message, No cache boundary
+    (3, None),  # 3 messages, No cache boundary
+    (5, 1),  # 5 messages, cache at index 2
+    (10, 6),  # 10 messages, cache at index 6
+])
+def test_caching_conversation_history(message_count, expected_cache_boundary):
+    """Test that conversation history is cached at the keep_recent boundary for different message counts."""
+    
+    # Create messages based on the parameter
+    messages: List[ChatCompletionMessageParam] = [
+        ChatCompletionSystemMessageParam(role="system", content="You are helpful.")
+    ]
+    
+    # Add message pairs
+    for i in range(message_count):
+        messages.append(ChatCompletionUserMessageParam(role="user", content=f"Message {i+1}"))
+    
+    system_prompt, anthropic_messages = get_anthropic_system_prompt_and_messages_with_caching(messages)
+    
+    # System prompt should have cache control
+    assert isinstance(system_prompt, list)
+    assert system_prompt[0]["cache_control"] == {"type": "ephemeral"}
+    
+    print(anthropic_messages)
+
+    if expected_cache_boundary is None:
+        # Verify no cache boundry 
+        assert all("cache_control" not in str(message) for message in anthropic_messages)
+    else:
+        # Other messages should not have cache control
+        for i, message in enumerate(anthropic_messages):
+            if i == expected_cache_boundary:
+                assert anthropic_messages[expected_cache_boundary]["content"][0]["cache_control"] == {"type": "ephemeral"}
+            else:
+                assert "cache_control" not in str(message)
+
+def test_caching_with_mixed_content():
+    """Test caching with mixed text and image content."""
+    messages: List[ChatCompletionMessageParam] = [
+        ChatCompletionSystemMessageParam(role="system", content="You are a helpful assistant."),
+        ChatCompletionUserMessageParam(role="user", content=[
+            {"type": "text", "text": "Here is an image:"},
+            {"type": "image_url", "image_url": {"url": DUMMY_IMAGE_DATA_URL}}
+        ])
+    ]
+    system_prompt, anthropic_messages = get_anthropic_system_prompt_and_messages_with_caching(messages)
+    
+    # System prompt should have cache control
+    assert isinstance(system_prompt, list)
+    assert system_prompt[0]["cache_control"] == {"type": "ephemeral"}
+    
+    # User message should NOT have cache control (only 1 message, so boundary is invalid)
+    user_message = anthropic_messages[0]
+    assert user_message["role"] == "user"
+    assert isinstance(user_message["content"], list)
+    assert len(user_message["content"]) == 2
+    
+    # No content blocks should have cache control (too few messages to cache)
+    assert user_message["content"][0]["type"] == "text"
+    assert "cache_control" not in user_message["content"][0]
+    assert user_message["content"][1]["type"] == "image"
+    assert "cache_control" not in user_message["content"][1]

mito_ai/tests/streamlit_conversion/test_apply_search_replace.py

@@ -0,0 +1,226 @@
+# Copyright (c) Saga Inc.
+# Distributed under the terms of the GNU Affero General Public License v3.0 License.
+
+from mito_ai.utils.error_classes import StreamlitConversionError
+import pytest
+from mito_ai.streamlit_conversion.search_replace_utils import apply_search_replace
+
+
+@pytest.mark.parametrize("original_text,search_replace_pairs,expected_result", [
+    # Test case 1: Simple title change
+    (
+        """import streamlit as st
+
+st.markdown(\"\"\"
+    <style>
+        #MainMenu {visibility: hidden;}
+        .stAppDeployButton {display:none;}
+        footer {visibility: hidden;}
+        .stMainBlockContainer {padding: 2rem 1rem 2rem 1rem;}
+    </style>
+\"\"\", unsafe_allow_html=True)
+
+st.title("Simple Calculation")
+
+x = 5
+y = 10
+result = x + y
+
+st.write(f"x = {x}")
+st.write(f"y = {y}")
+st.write(f"x + y = {result}")""",
+        [("st.title(\"Simple Calculation\")", "st.title(\"Math Examples\")")],
+        """import streamlit as st
+
+st.markdown(\"\"\"
+    <style>
+        #MainMenu {visibility: hidden;}
+        .stAppDeployButton {display:none;}
+        footer {visibility: hidden;}
+        .stMainBlockContainer {padding: 2rem 1rem 2rem 1rem;}
+    </style>
+\"\"\", unsafe_allow_html=True)
+
+st.title("Math Examples")
+
+x = 5
+y = 10
+result = x + y
+
+st.write(f"x = {x}")
+st.write(f"y = {y}")
+st.write(f"x + y = {result}")"""
+    ),
+    
+    # Test case 2: Add new content
+    (
+        """import streamlit as st
+
+st.title("My App")""",
+        [("st.title(\"My App\")", """st.title("My App")
+st.header("Welcome")
+st.write("This is a test app")""")],
+        """import streamlit as st
+
+st.title("My App")
+st.header("Welcome")
+st.write("This is a test app")"""
+    ),
+    
+    # Test case 3: Remove lines
+    (
+        """import streamlit as st
+
+st.header("Welcome")
+st.title("My App")
+st.write("This is a test app")""",
+        [("""st.header("Welcome")
+st.title("My App")
+st.write("This is a test app")""", "st.title(\"My App\")")],
+        """import streamlit as st
+
+st.title("My App")"""
+    ),
+    
+    # Test case 4: Multiple replacements
+    (
+        """import streamlit as st
+
+st.title("Old Title")
+x = 5
+y = 10
+st.write("Old message")""",
+        [
+            ("st.title(\"Old Title\")", "st.title(\"New Title\")"),
+            ("st.write(\"Old message\")", "st.write(\"New message\")")
+        ],
+        """import streamlit as st
+
+st.title("New Title")
+x = 5
+y = 10
+st.write("New message")"""
+    ),
+    
+    # Test case 5: Empty search/replace pairs
+    (
+        """import streamlit as st
+
+st.title("My App")""",
+        [],
+        """import streamlit as st
+
+st.title("My App")"""
+    ),
+    
+    # Test case 6: Complex replacement with context
+    (
+        """import streamlit as st
+
+# This is a comment
+st.title("Old Title")
+# Another comment
+x = 5
+y = 10
+# Final comment""",
+        [("""# This is a comment
+st.title("Old Title")
+# Another comment""", """# This is a comment
+st.title("New Title")
+# Another comment""")],
+        """import streamlit as st
+
+# This is a comment
+st.title("New Title")
+# Another comment
+x = 5
+y = 10
+# Final comment"""
+    ),
+    
+    # Test case 7: Replace multiple consecutive lines
+    (
+        """import streamlit as st
+
+st.title("My App")
+st.write("Line 1")
+st.write("Line 2")
+st.write("Line 3")
+
+x = 5""",
+        [("""st.write("Line 1")
+st.write("Line 2")
+st.write("Line 3")""", "st.write(\"New content\")")],
+        """import streamlit as st
+
+st.title("My App")
+st.write("New content")
+
+x = 5"""
+    ),
+    
+    # Test case 8: Add lines at the beginning
+    (
+        """import streamlit as st
+
+st.title("My App")""",
+        [("import streamlit as st", """import pandas as pd
+import streamlit as st""")],
+        """import pandas as pd
+import streamlit as st
+
+st.title("My App")"""
+    ),
+    
+    # Test case 9: Add lines at the end
+    (
+        """import streamlit as st
+
+st.title("My App")""",
+        [("st.title(\"My App\")", """st.title("My App")
+
+st.write("Footer content")
+st.write("More footer")""")],
+        """import streamlit as st
+
+st.title("My App")
+
+st.write("Footer content")
+st.write("More footer")"""
+    ),
+    
+    # Test case 10: Add emoji to streamlit app title
+    (
+        """import streamlit as st
+
+st.title("My App")
+st.write("Welcome to my application")""",
+        [("st.title(\"My App\")", "st.title(\"🚀 My App\")")],
+        """import streamlit as st
+
+st.title("🚀 My App")
+st.write("Welcome to my application")"""
+    )
+])
+def test_apply_search_replace(original_text, search_replace_pairs, expected_result):
+    """Test the apply_search_replace function with various search/replace scenarios."""
+    result = apply_search_replace(original_text, search_replace_pairs)
+    
+    print(f"Original text: {repr(original_text)}")
+    print(f"Search/replace pairs: {search_replace_pairs}")
+    print(f"Expected result: {repr(expected_result)}")
+    print(f"Result: {repr(result)}")
+    
+    assert result == expected_result
+
+
+def test_apply_search_replace_search_not_found():
+    """Test that ValueError is raised when search text is not found."""
+    with pytest.raises(StreamlitConversionError, match="Search text not found"):
+        apply_search_replace("st.title(\"My App\")", [("st.title(\"Not Found\")", "st.title(\"New Title\")")])
+
+
+def test_apply_search_replace_multiple_matches():
+    """Test that ValueError is raised when search text is found multiple times."""
+    with pytest.raises(StreamlitConversionError, match="Search text found 2 times"):
+        apply_search_replace("st.write(\"Hello\")\nst.write(\"Hello\")", [("st.write(\"Hello\")", "st.write(\"Hi\")")])