mito-ai 0.1.45__py3-none-any.whl → 0.1.46__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_patch_to_text.py

@@ -0,0 +1,368 @@
+# Copyright (c) Saga Inc.
+# Distributed under the terms of the GNU Affero General Public License v3.0 License.
+
+import pytest
+from mito_ai.streamlit_conversion.agent_utils import apply_patch_to_text
+
+
+@pytest.mark.parametrize("original_text,diff,expected_result", [
+    # Test case 1: Simple line replacement
+    (
+        """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}")""",
+        """--- a/app.py
++++ b/app.py
+@@ -10,4 +10,4 @@
+ \"\"\", unsafe_allow_html=True)
+
+-st.title("Simple Calculation")
++st.title("Math Examples")
+
+ x = 5""",
+        """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 lines
+    (
+        """import streamlit as st
+
+st.title("My App")""",
+        """--- a/app.py
++++ b/app.py
+@@ -1,3 +1,5 @@
+ import streamlit as st
+
++st.header("Welcome")
+ st.title("My App")
++st.write("This is a test app")""",
+        """import streamlit as st
+
+st.header("Welcome")
+st.title("My App")
+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")""",
+        """--- a/app.py
++++ b/app.py
+@@ -1,5 +1,3 @@
+ import streamlit as st
+
+-st.header("Welcome")
+ st.title("My App")
+-st.write("This is a test app")""",
+        """import streamlit as st
+
+st.title("My App")
+"""
+    ),
+    
+    # Test case 4: Single hunk with multiple changes
+    (
+        """import streamlit as st
+
+st.title("Old Title")
+x = 5
+y = 10
+st.write("Old message")""",
+        """--- a/app.py
++++ b/app.py
+@@ -1,6 +1,6 @@
+ import streamlit as st
+
+-st.title("Old Title")
++st.title("New Title")
+ x = 5
+ y = 10
+-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 diff
+    (
+        """import streamlit as st
+
+st.title("My App")""",
+        "",
+        """import streamlit as st
+
+st.title("My App")"""
+    ),
+    
+    # Test case 6: Whitespace-only changes
+    (
+        """import streamlit as st
+st.title("My App")""",
+        """--- a/app.py
++++ b/app.py
+@@ -1,2 +1,2 @@
+ import streamlit as st
+-st.title("My App")
++st.title("My App")""",
+        """import streamlit as st
+st.title("My App")"""
+    ),
+    
+    # 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""",
+        """--- a/app.py
++++ b/app.py
+@@ -3,4 +3,2 @@
+ st.title("My App")
+-st.write("Line 1")
+-st.write("Line 2")
+-st.write("Line 3")
++st.write("New content")
+
+ x = 5""",
+        """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")""",
+        """--- a/app.py
++++ b/app.py
+@@ -1,2 +1,3 @@
++import pandas as pd
+ import streamlit as st
+
+ st.title("My App")""",
+        """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")""",
+        """--- a/app.py
++++ b/app.py
+@@ -3,1 +3,4 @@
+ 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: Complex replacement with context
+    (
+        """import streamlit as st
+
+# This is a comment
+st.title("Old Title")
+# Another comment
+x = 5
+y = 10
+# Final comment""",
+        """--- a/app.py
++++ b/app.py
+@@ -2,4 +2,4 @@
+
+ # This is a comment
+-st.title("Old Title")
++st.title("New Title")
+ # Another comment
+ x = 5""",
+        """import streamlit as st
+
+# This is a comment
+st.title("New Title")
+# Another comment
+x = 5
+y = 10
+# Final comment"""
+    ),
+    
+    # Test case 11: Simple multiple hunks - title change and add footer
+    (
+        """import streamlit as st
+
+st.title("Old App")
+st.write("Hello World")""",
+        """--- a/app.py
++++ b/app.py
+@@ -3,1 +3,1 @@
+-st.title("Old App")
++st.title("New App")
+@@ -4,1 +4,3 @@
+ st.write("Hello World")
++
++st.write("Footer text")""",
+        """import streamlit as st
+
+st.title("New App")
+st.write("Hello World")
+
+st.write("Footer text")"""
+    ),
+    
+    # Test case 12: Simple multiple hunks - remove and add lines
+    (
+        """import streamlit as st
+st.write("Remove this")
+st.title("My App")""",
+        """--- a/app.py
++++ b/app.py
+@@ -2,1 +2,0 @@
+-st.write("Remove this")
+@@ -3,1 +3,2 @@
+ st.title("My App")
++st.write("Add this")""",
+        """import streamlit as st
+st.title("My App")
+st.write("Add this")"""
+    ),
+    
+    # Test case 13: Multiple hunks with context lines
+    (
+        """import streamlit as st
+
+st.title("Old Title")
+st.write("Some content")
+
+st.write("Old message")""",
+        """--- a/app.py
++++ b/app.py
+@@ -3,1 +3,1 @@
+-st.title("Old Title")
++st.title("New Title")
+@@ -6,1 +6,1 @@
+-st.write("Old message")
++st.write("New message")""",
+        """import streamlit as st
+
+st.title("New Title")
+st.write("Some content")
+
+st.write("New message")"""
+    ),
+    
+        # Test case 14: Multiple hunks - add imports and change content
+        (
+            """import streamlit as st
+
+st.title("My App")
+st.write("Hello")
+st.write("World")""",
+            """--- a/app.py
++++ b/app.py
+@@ -1,1 +1,3 @@
+ import streamlit as st
++import pandas as pd
++import numpy as np
+@@ -4,2 +4,2 @@
+ st.write("Hello")
+-st.write("World")
++st.write("World!")""",
+            """import streamlit as st
+import pandas as pd
+import numpy as np
+
+st.title("My App")
+st.write("Hello")
+st.write("World!")"""
+        ),
+        
+        # Test case 15: Add emoji to streamlit app title
+        (
+            """import streamlit as st
+
+st.title("My App")
+st.write("Welcome to my application")""",
+            """--- a/app.py
++++ b/app.py
+@@ -3,1 +3,1 @@
+-st.title("My App")
++st.title("🚀 My App")
+ 
+ st.write("Welcome to my application")""",
+            """import streamlit as st
+
+st.title("🚀 My App")
+st.write("Welcome to my application")"""
+        )
+])
+def test_apply_patch_to_text(original_text, diff, expected_result):
+    """Test the apply_patch_to_text function with various diff scenarios."""
+    result = apply_patch_to_text(original_text, diff)
+    
+    print(f"Original text: {repr(original_text)}")
+    print(f"Diff: {repr(diff)}")
+    print(f"Expected result: {repr(expected_result)}")
+    print(f"Result: {repr(result)}")
+    
+    assert result == expected_result