#### What this tests ####
#    This tests if prompts are being correctly formatted
import os
import sys

import pytest

sys.path.insert(0, os.path.abspath("../.."))

from typing import List

# from litellm.litellm_core_utils.prompt_templates.factory import prompt_factory
import litellm
from litellm.litellm_core_utils.prompt_templates.factory import (
    _bedrock_tools_pt,
    anthropic_messages_pt,
    anthropic_pt,
    claude_2_1_pt,
    convert_to_anthropic_image_obj,
    convert_to_anthropic_tool_invoke,
    convert_url_to_base64,
    create_anthropic_image_param,
    llama_2_chat_pt,
    prompt_factory,
)
from litellm.litellm_core_utils.prompt_templates.common_utils import (
    get_completion_messages,
)
from litellm.llms.vertex_ai.gemini.transformation import (
    _gemini_convert_messages_with_history,
)
from litellm.types.llms.openai import AllMessageValues
from unittest.mock import MagicMock, patch


def test_llama_3_prompt():
    messages = [
        {"role": "system", "content": "You are a good bot"},
        {"role": "user", "content": "Hey, how's it going?"},
    ]
    received_prompt = prompt_factory(
        model="meta-llama/Meta-Llama-3-8B-Instruct", messages=messages
    )
    print(f"received_prompt: {received_prompt}")

    expected_prompt = """<|begin_of_text|><|start_header_id|>system<|end_header_id|>\n\nYou are a good bot<|eot_id|><|start_header_id|>user<|end_header_id|>\n\nHey, how's it going?<|eot_id|><|start_header_id|>assistant<|end_header_id|>\n\n"""
    assert received_prompt == expected_prompt


def test_codellama_prompt_format():
    messages = [
        {"role": "system", "content": "You are a good bot"},
        {"role": "user", "content": "Hey, how's it going?"},
    ]
    expected_prompt = "<s>[INST] <<SYS>>\nYou are a good bot\n<</SYS>>\n [/INST]\n[INST] Hey, how's it going? [/INST]\n"
    assert llama_2_chat_pt(messages) == expected_prompt


def test_claude_2_1_pt_formatting():
    # Test case: User only, should add Assistant
    messages = [{"role": "user", "content": "Hello"}]
    expected_prompt = "\n\nHuman: Hello\n\nAssistant: "
    assert claude_2_1_pt(messages) == expected_prompt

    # Test case: System, User, and Assistant "pre-fill" sequence,
    #            Should return pre-fill
    messages = [
        {"role": "system", "content": "You are a helpful assistant."},
        {"role": "user", "content": 'Please return "Hello World" as a JSON object.'},
        {"role": "assistant", "content": "{"},
    ]
    expected_prompt = 'You are a helpful assistant.\n\nHuman: Please return "Hello World" as a JSON object.\n\nAssistant: {'
    assert claude_2_1_pt(messages) == expected_prompt

    # Test case: System, Assistant sequence, should insert blank Human message
    #            before Assistant pre-fill
    messages = [
        {"role": "system", "content": "You are a storyteller."},
        {"role": "assistant", "content": "Once upon a time, there "},
    ]
    expected_prompt = (
        "You are a storyteller.\n\nHuman: \n\nAssistant: Once upon a time, there "
    )
    assert claude_2_1_pt(messages) == expected_prompt

    # Test case: System, User sequence
    messages = [
        {"role": "system", "content": "System reboot"},
        {"role": "user", "content": "Is everything okay?"},
    ]
    expected_prompt = "System reboot\n\nHuman: Is everything okay?\n\nAssistant: "
    assert claude_2_1_pt(messages) == expected_prompt


def test_anthropic_pt_formatting():
    # Test case: User only, should add Assistant
    messages = [{"role": "user", "content": "Hello"}]
    expected_prompt = "\n\nHuman: Hello\n\nAssistant: "
    assert anthropic_pt(messages) == expected_prompt

    # Test case: System, User, and Assistant "pre-fill" sequence,
    #            Should return pre-fill
    messages = [
        {"role": "system", "content": "You are a helpful assistant."},
        {"role": "user", "content": 'Please return "Hello World" as a JSON object.'},
        {"role": "assistant", "content": "{"},
    ]
    expected_prompt = '\n\nHuman: <admin>You are a helpful assistant.</admin>\n\nHuman: Please return "Hello World" as a JSON object.\n\nAssistant: {'
    assert anthropic_pt(messages) == expected_prompt

    # Test case: System, Assistant sequence, should NOT insert blank Human message
    #            before Assistant pre-fill, because "System" messages are Human
    #            messages wrapped with <admin></admin>
    messages = [
        {"role": "system", "content": "You are a storyteller."},
        {"role": "assistant", "content": "Once upon a time, there "},
    ]
    expected_prompt = "\n\nHuman: <admin>You are a storyteller.</admin>\n\nAssistant: Once upon a time, there "
    assert anthropic_pt(messages) == expected_prompt

    # Test case: System, User sequence
    messages = [
        {"role": "system", "content": "System reboot"},
        {"role": "user", "content": "Is everything okay?"},
    ]
    expected_prompt = "\n\nHuman: <admin>System reboot</admin>\n\nHuman: Is everything okay?\n\nAssistant: "
    assert anthropic_pt(messages) == expected_prompt


def test_anthropic_messages_nested_pt():

    messages = [
        {"content": [{"text": "here is a task", "type": "text"}], "role": "user"},
        {
            "content": [{"text": "sure happy to help", "type": "text"}],
            "role": "assistant",
        },
        {
            "content": [
                {
                    "text": "Here is a screenshot of the current desktop with the "
                    "mouse coordinates (500, 350). Please select an action "
                    "from the provided schema.",
                    "type": "text",
                }
            ],
            "role": "user",
        },
    ]

    new_messages = anthropic_messages_pt(
        messages, model="claude-3-sonnet-20240229", llm_provider="anthropic"
    )

    assert isinstance(new_messages[1]["content"][0]["text"], str)


# codellama_prompt_format()
def test_bedrock_tool_calling_pt():
    tools = [
        {
            "type": "function",
            "function": {
                "name": "get_current_weather",
                "description": "Get the current weather in a given location",
                "parameters": {
                    "type": "object",
                    "properties": {
                        "location": {
                            "type": "string",
                            "description": "The city and state, e.g. San Francisco, CA",
                        },
                        "unit": {"type": "string", "enum": ["celsius", "fahrenheit"]},
                    },
                    "required": ["location"],
                },
            },
        }
    ]
    converted_tools = _bedrock_tools_pt(tools=tools)

    print(converted_tools)


def test_convert_url_to_img():
    response_url = convert_url_to_base64(
        url="https://images.pexels.com/photos/1319515/pexels-photo-1319515.jpeg?auto=compress&cs=tinysrgb&w=1260&h=750&dpr=1"
    )

    assert "image/jpeg" in response_url


@pytest.mark.parametrize(
    "url, expected_media_type",
    [
        ("data:image/jpeg;base64,1234", "image/jpeg"),
        ("data:application/pdf;base64,1234", "application/pdf"),
        (r"data:image\/jpeg;base64,1234", "image/jpeg"),
    ],
)
def test_base64_image_input(url, expected_media_type):
    response = convert_to_anthropic_image_obj(openai_image_url=url, format=None)

    assert response["media_type"] == expected_media_type


def test_create_anthropic_image_param_with_http_url():
    """Test that HTTP/HTTPS URLs are passed as URL references, not base64."""
    image_param = create_anthropic_image_param(
        "https://example.com/image.jpg", format=None
    )

    assert image_param["type"] == "image"
    assert image_param["source"]["type"] == "url"
    assert image_param["source"]["url"] == "https://example.com/image.jpg"


def test_create_anthropic_image_param_with_https_url():
    """Test that HTTPS URLs are passed as URL references."""
    image_param = create_anthropic_image_param(
        "https://example.com/image.png", format=None
    )

    assert image_param["type"] == "image"
    assert image_param["source"]["type"] == "url"
    assert image_param["source"]["url"] == "https://example.com/image.png"


def test_create_anthropic_image_param_with_dict_input():
    """Test that dict input with URL is handled correctly."""
    image_param = create_anthropic_image_param(
        {"url": "https://example.com/image.jpg", "format": "image/jpeg"}, format=None
    )

    assert image_param["type"] == "image"
    assert image_param["source"]["type"] == "url"
    assert image_param["source"]["url"] == "https://example.com/image.jpg"


def test_create_anthropic_image_param_with_base64_data_uri():
    """Test that data URIs are converted to base64."""
    image_param = create_anthropic_image_param(
        "data:image/jpeg;base64,/9j/4AAQSkZJRg==", format=None
    )

    assert image_param["type"] == "image"
    assert image_param["source"]["type"] == "base64"
    assert image_param["source"]["media_type"] == "image/jpeg"
    assert image_param["source"]["data"] == "/9j/4AAQSkZJRg=="


def test_create_anthropic_image_param_with_format_override():
    """Test that format parameter can override media type."""
    image_param = create_anthropic_image_param(
        "data:image/jpeg;base64,1234", format="image/png"
    )

    assert image_param["type"] == "image"
    assert image_param["source"]["type"] == "base64"
    assert image_param["source"]["media_type"] == "image/png"


def test_anthropic_messages_pt_with_url_image():
    """Test that anthropic_messages_pt correctly handles HTTP/HTTPS URLs as URL references."""
    messages = [
        {
            "role": "user",
            "content": [
                {"type": "text", "text": "What's in this image?"},
                {
                    "type": "image_url",
                    "image_url": "https://example.com/image.jpg",
                },
            ],
        }
    ]

    result = anthropic_messages_pt(
        messages=messages, model="claude-3-5-sonnet", llm_provider="anthropic"
    )

    assert len(result) == 1
    assert result[0]["role"] == "user"
    assert isinstance(result[0]["content"], list)
    assert len(result[0]["content"]) == 2

    # Check text content
    assert result[0]["content"][0]["type"] == "text"

    # Check image content - should be URL reference, not base64
    assert result[0]["content"][1]["type"] == "image"
    assert result[0]["content"][1]["source"]["type"] == "url"
    assert result[0]["content"][1]["source"]["url"] == "https://example.com/image.jpg"


def test_anthropic_messages_pt_with_base64_image():
    """Test that anthropic_messages_pt correctly handles data URIs as base64."""
    messages = [
        {
            "role": "user",
            "content": [
                {"type": "text", "text": "What's in this image?"},
                {
                    "type": "image_url",
                    "image_url": "data:image/jpeg;base64,/9j/4AAQSkZJRg==",
                },
            ],
        }
    ]

    result = anthropic_messages_pt(
        messages=messages, model="claude-3-5-sonnet", llm_provider="anthropic"
    )

    assert len(result) == 1
    assert result[0]["role"] == "user"
    assert isinstance(result[0]["content"], list)
    assert len(result[0]["content"]) == 2

    # Check image content - should be base64, not URL
    assert result[0]["content"][1]["type"] == "image"
    assert result[0]["content"][1]["source"]["type"] == "base64"
    assert result[0]["content"][1]["source"]["media_type"] == "image/jpeg"


def test_anthropic_messages_tool_call():
    messages = [
        {
            "role": "user",
            "content": "Would development of a software platform be under ASC 350-40 or ASC 985?",
        },
        {
            "role": "assistant",
            "content": "",
            "tool_call_id": "bc8cb4b6-88c4-4138-8993-3a9d9cd51656",
            "tool_calls": [
                {
                    "id": "bc8cb4b6-88c4-4138-8993-3a9d9cd51656",
                    "function": {
                        "arguments": '{"completed_steps": [], "next_steps": [{"tool_name": "AccountingResearchTool", "description": "Research ASC 350-40 to understand its scope and applicability to software development."}, {"tool_name": "AccountingResearchTool", "description": "Research ASC 985 to understand its scope and applicability to software development."}, {"tool_name": "AccountingResearchTool", "description": "Compare the scopes of ASC 350-40 and ASC 985 to determine which is more applicable to software platform development."}], "learnings": [], "potential_issues": ["The distinction between the two standards might not be clear-cut for all types of software development.", "There might be specific circumstances or details about the software platform that could affect which standard applies."], "missing_info": ["Specific details about the type of software platform being developed (e.g., for internal use or for sale).", "Whether the entity developing the software is also the end-user or if it\'s being developed for external customers."], "done": false, "required_formatting": null}',
                        "name": "TaskPlanningTool",
                    },
                    "type": "function",
                }
            ],
        },
        {
            "role": "function",
            "content": '{"completed_steps":[],"next_steps":[{"tool_name":"AccountingResearchTool","description":"Research ASC 350-40 to understand its scope and applicability to software development."},{"tool_name":"AccountingResearchTool","description":"Research ASC 985 to understand its scope and applicability to software development."},{"tool_name":"AccountingResearchTool","description":"Compare the scopes of ASC 350-40 and ASC 985 to determine which is more applicable to software platform development."}],"formatting_step":null}',
            "name": "TaskPlanningTool",
            "tool_call_id": "bc8cb4b6-88c4-4138-8993-3a9d9cd51656",
        },
    ]

    translated_messages = anthropic_messages_pt(
        messages, model="claude-3-sonnet-20240229", llm_provider="anthropic"
    )

    print(translated_messages)

    assert (
        translated_messages[-1]["content"][0]["tool_use_id"]
        == "bc8cb4b6-88c4-4138-8993-3a9d9cd51656"
    )


def test_anthropic_cache_controls_pt():
    "see anthropic docs for this: https://docs.anthropic.com/en/docs/build-with-claude/prompt-caching#continuing-a-multi-turn-conversation"
    messages = [
        # marked for caching with the cache_control parameter, so that this checkpoint can read from the previous cache.
        {
            "role": "user",
            "content": [
                {
                    "type": "text",
                    "text": "What are the key terms and conditions in this agreement?",
                    "cache_control": {"type": "ephemeral"},
                }
            ],
        },
        {
            "role": "assistant",
            "content": "Certainly! the key terms and conditions are the following: the contract is 1 year long for $10/mo",
        },
        # The final turn is marked with cache-control, for continuing in followups.
        {
            "role": "user",
            "content": [
                {
                    "type": "text",
                    "text": "What are the key terms and conditions in this agreement?",
                    "cache_control": {"type": "ephemeral"},
                }
            ],
        },
        {
            "role": "assistant",
            "content": "Certainly! the key terms and conditions are the following: the contract is 1 year long for $10/mo",
            "cache_control": {"type": "ephemeral"},
        },
    ]

    translated_messages = anthropic_messages_pt(
        messages, model="claude-3-5-sonnet-20240620", llm_provider="anthropic"
    )

    for i, msg in enumerate(translated_messages):
        if i == 0:
            assert msg["content"][0]["cache_control"] == {"type": "ephemeral"}
        elif i == 1:
            assert "cache_controls" not in msg["content"][0]
        elif i == 2:
            assert msg["content"][0]["cache_control"] == {"type": "ephemeral"}
        elif i == 3:
            assert msg["content"][0]["cache_control"] == {"type": "ephemeral"}

    print("translated_messages: ", translated_messages)


def test_anthropic_cache_controls_tool_calls_pt():
    """
    Tests that cache_control is properly set in tool_calls when converting messages
    for the Anthropic API.
    """
    messages = [
        {
            "role": "user",
            "content": "Can you help me get the weather?",
        },
        {
            "role": "assistant",
            "content": "",
            "tool_calls": [
                {
                    "id": "weather-tool-id-123",
                    "function": {
                        "arguments": '{"location": "San Francisco"}',
                        "name": "get_weather",
                    },
                    "type": "function",
                }
            ],
            "cache_control": {"type": "ephemeral"},
        },
        {
            "role": "function",
            "content": '{"temperature": 72, "unit": "fahrenheit", "description": "sunny"}',
            "name": "get_weather",
            "tool_call_id": "weather-tool-id-123",
            "cache_control": {"type": "ephemeral"},
        },
    ]

    translated_messages = anthropic_messages_pt(
        messages, model="claude-3-sonnet-20240229", llm_provider="anthropic"
    )

    print("Translated tool call messages:", translated_messages)

    assert translated_messages[0]["role"] == "user"

    assert translated_messages[1]["role"] == "assistant"
    for content_item in translated_messages[1]["content"]:
        if content_item["type"] == "tool_use":
            assert "cache_control" not in content_item
            assert content_item["name"] == "get_weather"

    assert translated_messages[2]["role"] == "user"
    for content_item in translated_messages[2]["content"]:
        if content_item["type"] == "tool_result":
            assert content_item["cache_control"] == {"type": "ephemeral"}


@pytest.mark.parametrize("provider", ["bedrock", "anthropic"])
def test_bedrock_parallel_tool_calling_pt(provider):
    """
    Make sure parallel tool call blocks are merged correctly - https://github.com/BerriAI/litellm/issues/5277
    """
    from litellm.litellm_core_utils.prompt_templates.factory import (
        _bedrock_converse_messages_pt,
    )
    from litellm.types.utils import ChatCompletionMessageToolCall, Function, Message

    messages = [
        {
            "role": "user",
            "content": "What's the weather like in San Francisco, Tokyo, and Paris? - give me 3 responses",
        },
        Message(
            content="Here are the current weather conditions for San Francisco, Tokyo, and Paris:",
            role="assistant",
            tool_calls=[
                ChatCompletionMessageToolCall(
                    index=1,
                    function=Function(
                        arguments='{"city": "New York"}',
                        name="get_current_weather",
                    ),
                    id="tooluse_XcqEBfm8R-2YVaPhDUHsPQ",
                    type="function",
                ),
                ChatCompletionMessageToolCall(
                    index=2,
                    function=Function(
                        arguments='{"city": "London"}',
                        name="get_current_weather",
                    ),
                    id="tooluse_VB9nk7UGRniVzGcaj6xrAQ",
                    type="function",
                ),
            ],
            function_call=None,
        ),
        {
            "tool_call_id": "tooluse_XcqEBfm8R-2YVaPhDUHsPQ",
            "role": "tool",
            "name": "get_current_weather",
            "content": "25 degrees celsius.",
        },
        {
            "tool_call_id": "tooluse_VB9nk7UGRniVzGcaj6xrAQ",
            "role": "tool",
            "name": "get_current_weather",
            "content": "28 degrees celsius.",
        },
    ]

    if provider == "bedrock":
        translated_messages = _bedrock_converse_messages_pt(
            messages=messages,
            model="anthropic.claude-3-sonnet-20240229-v1:0",
            llm_provider="bedrock",
        )
    else:
        translated_messages = anthropic_messages_pt(
            messages=messages,
            model="claude-3-sonnet-20240229-v1:0",
            llm_provider=provider,
        )
    print(translated_messages)

    number_of_messages = len(translated_messages)

    # assert last 2 messages are not the same role
    assert (
        translated_messages[number_of_messages - 1]["role"]
        != translated_messages[number_of_messages - 2]["role"]
    )


def test_vertex_only_image_user_message():
    base64_image = "/9j/2wCEAAgGBgcGBQ"

    messages = [
        {
            "role": "user",
            "content": [
                {
                    "type": "image_url",
                    "image_url": {"url": f"data:image/jpeg;base64,{base64_image}"},
                },
            ],
        },
    ]

    response = _gemini_convert_messages_with_history(
        messages=messages, model="gemini-1.5-pro"
    )

    expected_response = [
        {
            "role": "user",
            "parts": [
                {
                    "inline_data": {
                        "data": "/9j/2wCEAAgGBgcGBQ",
                        "mime_type": "image/jpeg",
                    }
                },
                {"text": " "},
            ],
        }
    ]

    assert len(response) == len(expected_response)
    for idx, content in enumerate(response):
        assert (
            content == expected_response[idx]
        ), "Invalid gemini input. Got={}, Expected={}".format(
            content, expected_response[idx]
        )


def test_no_messages_yields_user_text():
    """
    Test that contents are not empty and have text when called without messages
    This is to support blha blah
    """
    messages: List[AllMessageValues] = []

    contents = _gemini_convert_messages_with_history(messages=messages)

    expected_output = [{"role": "user", "parts": [{"text": " "}]}]

    assert contents == expected_output


def test_convert_url():
    convert_url_to_base64("https://picsum.photos/id/237/200/300")


def test_azure_tool_call_invoke_helper():
    messages = [
        {"role": "system", "content": "You are a helpful assistant."},
        {"role": "user", "content": "What is the weather in Copenhagen?"},
        {"role": "assistant", "function_call": {"name": "get_weather"}},
    ]

    transformed_messages = litellm.AzureOpenAIConfig().transform_request(
        model="gpt-4o",
        messages=messages,
        optional_params={},
        litellm_params={},
        headers={},
    )

    assert transformed_messages["messages"] == [
        {"role": "system", "content": "You are a helpful assistant."},
        {"role": "user", "content": "What is the weather in Copenhagen?"},
        {
            "role": "assistant",
            "function_call": {"name": "get_weather", "arguments": ""},
        },
    ]


@pytest.mark.parametrize(
    "messages, expected_messages, user_continue_message, assistant_continue_message",
    [
        (
            [
                {"role": "user", "content": "Hello!"},
                {"role": "assistant", "content": "Hello! How can I assist you today?"},
                {"role": "user", "content": "What is Databricks?"},
                {"role": "user", "content": "What is Azure?"},
                {"role": "assistant", "content": "I don't know anyything, do you?"},
            ],
            [
                {"role": "user", "content": "Hello!"},
                {
                    "role": "assistant",
                    "content": "Hello! How can I assist you today?",
                },
                {"role": "user", "content": "What is Databricks?"},
                {
                    "role": "assistant",
                    "content": "Please continue.",
                },
                {"role": "user", "content": "What is Azure?"},
                {
                    "role": "assistant",
                    "content": "I don't know anyything, do you?",
                },
                {
                    "role": "user",
                    "content": "Please continue.",
                },
            ],
            None,
            None,
        ),
        (
            [
                {"role": "user", "content": "Hello!"},
            ],
            [
                {"role": "user", "content": "Hello!"},
            ],
            None,
            None,
        ),
        (
            [
                {"role": "user", "content": "Hello!"},
                {"role": "user", "content": "What is Databricks?"},
            ],
            [
                {"role": "user", "content": "Hello!"},
                {"role": "assistant", "content": "Please continue."},
                {"role": "user", "content": "What is Databricks?"},
            ],
            None,
            None,
        ),
        (
            [
                {"role": "user", "content": "Hello!"},
                {"role": "user", "content": "What is Databricks?"},
                {"role": "user", "content": "What is Azure?"},
            ],
            [
                {"role": "user", "content": "Hello!"},
                {"role": "assistant", "content": "Please continue."},
                {"role": "user", "content": "What is Databricks?"},
                {
                    "role": "assistant",
                    "content": "Please continue.",
                },
                {"role": "user", "content": "What is Azure?"},
            ],
            None,
            None,
        ),
        (
            [
                {"role": "user", "content": "Hello!"},
                {
                    "role": "assistant",
                    "content": "Hello! How can I assist you today?",
                },
                {"role": "user", "content": "What is Databricks?"},
                {"role": "user", "content": "What is Azure?"},
                {"role": "assistant", "content": "I don't know anyything, do you?"},
                {"role": "assistant", "content": "I can't repeat sentences."},
            ],
            [
                {"role": "user", "content": "Hello!"},
                {
                    "role": "assistant",
                    "content": "Hello! How can I assist you today?",
                },
                {"role": "user", "content": "What is Databricks?"},
                {
                    "role": "assistant",
                    "content": "Please continue",
                },
                {"role": "user", "content": "What is Azure?"},
                {
                    "role": "assistant",
                    "content": "I don't know anyything, do you?",
                },
                {
                    "role": "user",
                    "content": "Ok",
                },
                {
                    "role": "assistant",
                    "content": "I can't repeat sentences.",
                },
                {"role": "user", "content": "Ok"},
            ],
            {
                "role": "user",
                "content": "Ok",
            },
            {
                "role": "assistant",
                "content": "Please continue",
            },
        ),
    ],
)
def test_ensure_alternating_roles(
    messages, expected_messages, user_continue_message, assistant_continue_message
):
    messages = get_completion_messages(
        messages=messages,
        assistant_continue_message=assistant_continue_message,
        user_continue_message=user_continue_message,
        ensure_alternating_roles=True,
    )

    print(messages)

    assert messages == expected_messages


def test_ensure_alternating_roles_with_tool_calls():
    """Fixes Regression in #18685"""
    messages = [
        {"role": "user", "content": "What's the weather?"},
        {
            "role": "assistant",
            "content": None,
            "tool_calls": [
                {
                    "id": "call_123",
                    "type": "function",
                    "function": {
                        "name": "get_weather",
                        "arguments": '{"location": "NYC"}',
                    },
                }
            ],
        },
        {"role": "tool", "tool_call_id": "call_123", "content": "72F, sunny"},
        {"role": "assistant", "content": "It's 72F and sunny in NYC."},
        {"role": "user", "content": "What about tomorrow?"},
        {"role": "user", "content": "And the day after?"},
        {"role": "user", "content": "What about next week?"},
    ]

    transformed_messages = get_completion_messages(
        messages=messages,
        assistant_continue_message=None,
        user_continue_message=None,
        ensure_alternating_roles=True,
    )

    assert transformed_messages == [
        {"role": "user", "content": "What's the weather?"},
        {
            "role": "assistant",
            "content": None,
            "tool_calls": [
                {
                    "id": "call_123",
                    "type": "function",
                    "function": {
                        "name": "get_weather",
                        "arguments": '{"location": "NYC"}',
                    },
                }
            ],
        },
        {"role": "tool", "tool_call_id": "call_123", "content": "72F, sunny"},
        {"role": "assistant", "content": "It's 72F and sunny in NYC."},
        {"role": "user", "content": "What about tomorrow?"},
        {"role": "assistant", "content": "Please continue."},
        {"role": "user", "content": "And the day after?"},
        {"role": "assistant", "content": "Please continue."},
        {"role": "user", "content": "What about next week?"},
    ]


def test_ensure_alternating_roles_three_consecutive_assistants():
    messages = [
        {"role": "user", "content": "Hello"},
        {"role": "assistant", "content": "A1"},
        {"role": "assistant", "content": "A2"},
        {"role": "assistant", "content": "A3"},
    ]

    transformed_messages = get_completion_messages(
        messages=messages,
        assistant_continue_message=None,
        user_continue_message=None,
        ensure_alternating_roles=True,
    )

    assert transformed_messages == [
        {"role": "user", "content": "Hello"},
        {"role": "assistant", "content": "A1"},
        {"role": "user", "content": "Please continue."},
        {"role": "assistant", "content": "A2"},
        {"role": "user", "content": "Please continue."},
        {"role": "assistant", "content": "A3"},
        {"role": "user", "content": "Please continue."},
    ]


def test_ensure_alternating_roles_inserts_assistant_continue_across_tool_chain():
    """[user, assistant(tc), tool, user] gets assistant_continue before the second user."""
    messages = [
        {"role": "user", "content": "Search for X"},
        {
            "role": "assistant",
            "content": None,
            "tool_calls": [
                {
                    "id": "c1",
                    "type": "function",
                    "function": {"name": "search", "arguments": "{}"},
                }
            ],
        },
        {"role": "tool", "tool_call_id": "c1", "content": "results"},
        {"role": "user", "content": "Thanks, now do Y"},
    ]

    transformed_messages = get_completion_messages(
        messages=messages,
        assistant_continue_message=None,
        user_continue_message=None,
        ensure_alternating_roles=True,
    )

    assert transformed_messages == [
        {"role": "user", "content": "Search for X"},
        {
            "role": "assistant",
            "content": None,
            "tool_calls": [
                {
                    "id": "c1",
                    "type": "function",
                    "function": {"name": "search", "arguments": "{}"},
                }
            ],
        },
        {"role": "tool", "tool_call_id": "c1", "content": "results"},
        {"role": "assistant", "content": "Please continue."},
        {"role": "user", "content": "Thanks, now do Y"},
    ]


def test_ensure_alternating_roles_assistant_tool_call_then_assistant():
    """
    Malformed [assistant(tc), assistant(no-tc), user]:
    user_continue inserts break between adjacents, then assistant_continue
    fills the counted-sequence gap.
    """
    messages = [
        {
            "role": "assistant",
            "content": None,
            "tool_calls": [
                {
                    "id": "c1",
                    "type": "function",
                    "function": {"name": "search", "arguments": "{}"},
                }
            ],
        },
        {"role": "assistant", "content": "Here's what I found."},
        {"role": "user", "content": "Thanks"},
    ]

    transformed_messages = get_completion_messages(
        messages=messages,
        assistant_continue_message=None,
        user_continue_message=None,
        ensure_alternating_roles=True,
    )

    assert transformed_messages == [
        {"role": "user", "content": "Please continue."},
        {
            "role": "assistant",
            "content": None,
            "tool_calls": [
                {
                    "id": "c1",
                    "type": "function",
                    "function": {"name": "search", "arguments": "{}"},
                }
            ],
        },
        {"role": "assistant", "content": "Please continue."},
        {"role": "user", "content": "Please continue."},
        {"role": "assistant", "content": "Here's what I found."},
        {"role": "user", "content": "Thanks"},
    ]


def test_ensure_alternating_roles_trailing_tool_call_assistant():
    messages = [
        {"role": "user", "content": "What's the weather?"},
        {
            "role": "assistant",
            "content": None,
            "tool_calls": [
                {
                    "id": "call_abc",
                    "type": "function",
                    "function": {
                        "name": "get_weather",
                        "arguments": '{"location": "NYC"}',
                    },
                }
            ],
        },
    ]

    transformed_messages = get_completion_messages(
        messages=messages,
        assistant_continue_message=None,
        user_continue_message=None,
        ensure_alternating_roles=True,
    )

    assert transformed_messages == [
        {"role": "user", "content": "What's the weather?"},
        {
            "role": "assistant",
            "content": None,
            "tool_calls": [
                {
                    "id": "call_abc",
                    "type": "function",
                    "function": {
                        "name": "get_weather",
                        "arguments": '{"location": "NYC"}',
                    },
                }
            ],
        },
        {"role": "assistant", "content": "Please continue."},
        {"role": "user", "content": "Please continue."},
    ]


def test_ensure_alternating_roles_multiple_tool_results():
    """[user, assistant(tc), tool, tool, user] — multiple tool results before next user."""
    messages = [
        {"role": "user", "content": "Search for X and Y"},
        {
            "role": "assistant",
            "content": None,
            "tool_calls": [
                {
                    "id": "c1",
                    "type": "function",
                    "function": {"name": "search_x", "arguments": "{}"},
                },
                {
                    "id": "c2",
                    "type": "function",
                    "function": {"name": "search_y", "arguments": "{}"},
                },
            ],
        },
        {"role": "tool", "tool_call_id": "c1", "content": "result X"},
        {"role": "tool", "tool_call_id": "c2", "content": "result Y"},
        {"role": "user", "content": "Thanks"},
    ]

    transformed_messages = get_completion_messages(
        messages=messages,
        assistant_continue_message=None,
        user_continue_message=None,
        ensure_alternating_roles=True,
    )

    assert transformed_messages == [
        {"role": "user", "content": "Search for X and Y"},
        {
            "role": "assistant",
            "content": None,
            "tool_calls": [
                {
                    "id": "c1",
                    "type": "function",
                    "function": {"name": "search_x", "arguments": "{}"},
                },
                {
                    "id": "c2",
                    "type": "function",
                    "function": {"name": "search_y", "arguments": "{}"},
                },
            ],
        },
        {"role": "tool", "tool_call_id": "c1", "content": "result X"},
        {"role": "tool", "tool_call_id": "c2", "content": "result Y"},
        {"role": "assistant", "content": "Please continue."},
        {"role": "user", "content": "Thanks"},
    ]


def test_ensure_alternating_roles_chained_tool_calls():
    """[user, assistant(tc), tool, assistant(tc), tool, user] — chained tool calls."""
    messages = [
        {"role": "user", "content": "Do multi-step task"},
        {
            "role": "assistant",
            "content": None,
            "tool_calls": [
                {
                    "id": "c1",
                    "type": "function",
                    "function": {"name": "step1", "arguments": "{}"},
                },
            ],
        },
        {"role": "tool", "tool_call_id": "c1", "content": "step1 done"},
        {
            "role": "assistant",
            "content": None,
            "tool_calls": [
                {
                    "id": "c2",
                    "type": "function",
                    "function": {"name": "step2", "arguments": "{}"},
                },
            ],
        },
        {"role": "tool", "tool_call_id": "c2", "content": "step2 done"},
        {"role": "user", "content": "What happened?"},
    ]

    transformed_messages = get_completion_messages(
        messages=messages,
        assistant_continue_message=None,
        user_continue_message=None,
        ensure_alternating_roles=True,
    )

    assert transformed_messages == [
        {"role": "user", "content": "Do multi-step task"},
        {
            "role": "assistant",
            "content": None,
            "tool_calls": [
                {
                    "id": "c1",
                    "type": "function",
                    "function": {"name": "step1", "arguments": "{}"},
                },
            ],
        },
        {"role": "tool", "tool_call_id": "c1", "content": "step1 done"},
        {
            "role": "assistant",
            "content": None,
            "tool_calls": [
                {
                    "id": "c2",
                    "type": "function",
                    "function": {"name": "step2", "arguments": "{}"},
                },
            ],
        },
        {"role": "tool", "tool_call_id": "c2", "content": "step2 done"},
        {"role": "assistant", "content": "Please continue."},
        {"role": "user", "content": "What happened?"},
    ]


def test_ensure_alternating_roles_system_prefix_with_tool_chain():
    """[system, user, assistant(tc), tool, user] — system prefix doesn't interfere."""
    messages = [
        {"role": "system", "content": "You are helpful."},
        {"role": "user", "content": "Search for X"},
        {
            "role": "assistant",
            "content": None,
            "tool_calls": [
                {
                    "id": "c1",
                    "type": "function",
                    "function": {"name": "search", "arguments": "{}"},
                },
            ],
        },
        {"role": "tool", "tool_call_id": "c1", "content": "results"},
        {"role": "user", "content": "Thanks"},
    ]

    transformed_messages = get_completion_messages(
        messages=messages,
        assistant_continue_message=None,
        user_continue_message=None,
        ensure_alternating_roles=True,
    )

    assert transformed_messages == [
        {"role": "system", "content": "You are helpful."},
        {"role": "user", "content": "Search for X"},
        {
            "role": "assistant",
            "content": None,
            "tool_calls": [
                {
                    "id": "c1",
                    "type": "function",
                    "function": {"name": "search", "arguments": "{}"},
                },
            ],
        },
        {"role": "tool", "tool_call_id": "c1", "content": "results"},
        {"role": "assistant", "content": "Please continue."},
        {"role": "user", "content": "Thanks"},
    ]


def test_alternating_roles_e2e():
    from litellm.llms.custom_httpx.http_handler import HTTPHandler
    import json

    litellm.set_verbose = True
    http_handler = HTTPHandler()

    with patch.object(http_handler, "post", new=MagicMock()) as mock_post:
        try:
            response = litellm.completion(
                **{
                    "model": "databricks/databricks-meta-llama-3-1-70b-instruct",
                    "messages": [
                        {"role": "user", "content": "Hello!"},
                        {
                            "role": "assistant",
                            "content": "Hello! How can I assist you today?",
                        },
                        {"role": "user", "content": "What is Databricks?"},
                        {"role": "user", "content": "What is Azure?"},
                        {
                            "role": "assistant",
                            "content": "I don't know anyything, do you?",
                        },
                        {"role": "assistant", "content": "I can't repeat sentences."},
                    ],
                    "user_continue_message": {
                        "role": "user",
                        "content": "Ok",
                    },
                    "assistant_continue_message": {
                        "role": "assistant",
                        "content": "Please continue",
                    },
                    "ensure_alternating_roles": True,
                },
                client=http_handler,
            )
        except Exception as e:
            print(f"error: {e}")

        assert mock_post.call_args.kwargs["data"] == json.dumps(
            {
                "model": "databricks-meta-llama-3-1-70b-instruct",
                "messages": [
                    {"role": "user", "content": "Hello!"},
                    {
                        "role": "assistant",
                        "content": "Hello! How can I assist you today?",
                    },
                    {"role": "user", "content": "What is Databricks?"},
                    {
                        "role": "assistant",
                        "content": "Please continue",
                    },
                    {"role": "user", "content": "What is Azure?"},
                    {
                        "role": "assistant",
                        "content": "I don't know anyything, do you?",
                    },
                    {
                        "role": "user",
                        "content": "Ok",
                    },
                    {
                        "role": "assistant",
                        "content": "I can't repeat sentences.",
                    },
                    {
                        "role": "user",
                        "content": "Ok",
                    },
                ],
            }
        )


def test_just_system_message():
    from litellm.litellm_core_utils.prompt_templates.factory import (
        _bedrock_converse_messages_pt,
    )

    with pytest.raises(litellm.BadRequestError) as e:
        _bedrock_converse_messages_pt(
            messages=[],
            model="anthropic.claude-3-sonnet-20240229-v1:0",
            llm_provider="bedrock",
        )
        assert "bedrock requires at least one non-system message" in str(e.value)


def test_convert_generic_image_chunk_to_openai_image_obj():
    from litellm.litellm_core_utils.prompt_templates.factory import (
        convert_generic_image_chunk_to_openai_image_obj,
        convert_to_anthropic_image_obj,
    )

    url = "https://i.pinimg.com/736x/b4/b1/be/b4b1becad04d03a9071db2817fc9fe77.jpg"
    image_obj = convert_to_anthropic_image_obj(url, format=None)
    url_str = convert_generic_image_chunk_to_openai_image_obj(image_obj)
    image_obj = convert_to_anthropic_image_obj(url_str, format=None)
    print(image_obj)


def test_hf_chat_template():
    from litellm.litellm_core_utils.prompt_templates.factory import (
        hf_chat_template,
    )

    model = "llama/arn:aws:bedrock:us-east-1:1234:imported-model/45d34re"
    litellm.register_prompt_template(
        model=model,
        tokenizer_config={
            "add_bos_token": True,
            "add_eos_token": False,
            "bos_token": {
                "__type": "AddedToken",
                "content": "",
                "lstrip": False,
                "normalized": True,
                "rstrip": False,
                "single_word": False,
            },
            "clean_up_tokenization_spaces": False,
            "eos_token": {
                "__type": "AddedToken",
                "content": "",
                "lstrip": False,
                "normalized": True,
                "rstrip": False,
                "single_word": False,
            },
            "legacy": True,
            "model_max_length": 16384,
            "pad_token": {
                "__type": "AddedToken",
                "content": "",
                "lstrip": False,
                "normalized": True,
                "rstrip": False,
                "single_word": False,
            },
            "sp_model_kwargs": {},
            "unk_token": None,
            "tokenizer_class": "LlamaTokenizerFast",
            "chat_template": "{% if not add_generation_prompt is defined %}{% set add_generation_prompt = false %}{% endif %}{% set ns = namespace(is_first=false, is_tool=false, is_output_first=true, system_prompt='') %}{%- for message in messages %}{%- if message['role'] == 'system' %}{% set ns.system_prompt = message['content'] %}{%- endif %}{%- endfor %}{{bos_token}}{{ns.system_prompt}}{%- for message in messages %}{%- if message['role'] == 'user' %}{%- set ns.is_tool = false -%}{{' ' + message['content']}}{%- endif %}{%- if message['role'] == 'assistant' and message['content'] is none %}{%- set ns.is_tool = false -%}{%- for tool in message['tool_calls']%}{%- if not ns.is_first %}{{' ' + tool['type'] + ' ' + tool['function']['name'] + '\n' + '```json' + '\n' + tool['function']['arguments'] + '\n' + '```' + ' '}}{%- set ns.is_first = true -%}{%- else %}{{' ' + tool['type'] + ' ' + tool['function']['name'] + '\n' + '```json' + '\n' + tool['function']['arguments'] + '\n' + '```' + ' '}}{{' '}}{%- endif %}{%- endfor %}{%- endif %}{%- if message['role'] == 'assistant' and message['content'] is not none %}{%- if ns.is_tool %}{{' ' + message['content'] + ' '}}{%- set ns.is_tool = false -%}{%- else %}{% set content = message['content'] %}{% if '</think>' in content %}{% set content = content.split('</think>')[-1] %}{% endif %}{{' ' + content + ' '}}{%- endif %}{%- endif %}{%- if message['role'] == 'tool' %}{%- set ns.is_tool = true -%}{%- if ns.is_output_first %}{{' ' + message['content'] + ' '}}{%- set ns.is_output_first = false %}{%- else %}{{' ' + message['content'] + ' '}}{%- endif %}{%- endif %}{%- endfor -%}{% if ns.is_tool %}{{' '}}{% endif %}{% if add_generation_prompt and not ns.is_tool %}{{' '}}{% endif %}",
        },
    )

    messages = [
        {"role": "system", "content": "You are a helpful assistant."},
        {"role": "user", "content": "What is the weather in Copenhagen?"},
    ]
    chat_template = hf_chat_template(model=model, messages=messages)
    print(chat_template)
    assert (
        chat_template.rstrip()
        == "You are a helpful assistant. What is the weather in Copenhagen?"
    )


def test_ollama_pt():
    from litellm.litellm_core_utils.prompt_templates.factory import ollama_pt

    messages = [
        {"role": "system", "content": "You are a helpful assistant."},
        {"role": "user", "content": "Hello!"},
    ]
    prompt = ollama_pt(model="ollama/llama3.1", messages=messages)
    print(prompt)


# ============ Server Tool Use Reconstruction Tests ============
# Fixes: https://github.com/BerriAI/litellm/issues/17737


def test_convert_to_anthropic_tool_invoke_regular_tool():
    """Test that regular tool_use is converted correctly."""
    tool_calls = [
        {
            "id": "toolu_01ABC123",
            "type": "function",
            "function": {
                "name": "get_weather",
                "arguments": '{"location": "San Francisco"}',
            },
        }
    ]

    result = convert_to_anthropic_tool_invoke(tool_calls)

    assert len(result) == 1
    assert result[0]["type"] == "tool_use"
    assert result[0]["id"] == "toolu_01ABC123"
    assert result[0]["name"] == "get_weather"
    assert result[0]["input"] == {"location": "San Francisco"}


def test_convert_to_anthropic_tool_invoke_sanitizes_invalid_ids():
    """Test that tool_use IDs with invalid characters are sanitized.

    Anthropic requires tool_use_id to match ^[a-zA-Z0-9_-]+$.
    IDs from external frameworks (e.g. MiniMax) may contain characters
    like colons that violate this pattern.
    """
    tool_calls = [
        {
            "id": "sessions_history:183",
            "type": "function",
            "function": {
                "name": "get_weather",
                "arguments": '{"location": "Boston"}',
            },
        },
        {
            "id": "composio.NOTION_SEARCH",
            "type": "function",
            "function": {
                "name": "search_notes",
                "arguments": '{"query": "test"}',
            },
        },
    ]

    result = convert_to_anthropic_tool_invoke(tool_calls)

    assert len(result) == 2
    # Colons replaced with underscores
    assert result[0]["id"] == "sessions_history_183"
    # Dots replaced with underscores
    assert result[1]["id"] == "composio_NOTION_SEARCH"
    # Valid IDs should pass through unchanged
    valid_tool_calls = [
        {
            "id": "toolu_01ABC-xyz_123",
            "type": "function",
            "function": {
                "name": "get_weather",
                "arguments": '{"location": "NYC"}',
            },
        }
    ]
    valid_result = convert_to_anthropic_tool_invoke(valid_tool_calls)
    assert valid_result[0]["id"] == "toolu_01ABC-xyz_123"


def test_convert_to_anthropic_tool_invoke_server_tool():
    """
    Test that server_tool_use (srvtoolu_) is reconstructed as server_tool_use.

    Fixes: https://github.com/BerriAI/litellm/issues/17737
    """
    tool_calls = [
        {
            "id": "srvtoolu_01ABC123",
            "type": "function",
            "function": {
                "name": "web_search",
                "arguments": '{"query": "elephant weight"}',
            },
        }
    ]

    result = convert_to_anthropic_tool_invoke(tool_calls)

    assert len(result) == 1
    assert result[0]["type"] == "server_tool_use"  # NOT tool_use
    assert result[0]["id"] == "srvtoolu_01ABC123"
    assert result[0]["name"] == "web_search"
    assert result[0]["input"] == {"query": "elephant weight"}


def test_convert_to_anthropic_tool_invoke_with_web_search_results():
    """
    Test that web_search_tool_result is included after server_tool_use.

    Fixes: https://github.com/BerriAI/litellm/issues/17737
    """
    tool_calls = [
        {
            "id": "srvtoolu_01ABC123",
            "type": "function",
            "function": {
                "name": "web_search",
                "arguments": '{"query": "elephant weight"}',
            },
        }
    ]

    web_search_results = [
        {
            "type": "web_search_tool_result",
            "tool_use_id": "srvtoolu_01ABC123",
            "content": [
                {
                    "type": "web_search_result",
                    "url": "https://example.com",
                    "title": "Elephant Facts",
                    "snippet": "Elephants weigh 5000 kg",
                }
            ],
        }
    ]

    result = convert_to_anthropic_tool_invoke(
        tool_calls, web_search_results=web_search_results
    )

    assert len(result) == 2
    # First: server_tool_use
    assert result[0]["type"] == "server_tool_use"
    assert result[0]["id"] == "srvtoolu_01ABC123"
    # Second: web_search_tool_result
    assert result[1]["type"] == "web_search_tool_result"
    assert result[1]["tool_use_id"] == "srvtoolu_01ABC123"


def test_convert_to_anthropic_tool_invoke_mixed_tools():
    """
    Test that mixed server and regular tools are reconstructed correctly.

    Fixes: https://github.com/BerriAI/litellm/issues/17737
    """
    tool_calls = [
        {
            "id": "srvtoolu_01ABC123",
            "type": "function",
            "function": {
                "name": "web_search",
                "arguments": '{"query": "elephant weight"}',
            },
        },
        {
            "id": "toolu_01XYZ789",
            "type": "function",
            "function": {"name": "add_numbers", "arguments": '{"a": 5000, "b": 100}'},
        },
    ]

    web_search_results = [
        {
            "type": "web_search_tool_result",
            "tool_use_id": "srvtoolu_01ABC123",
            "content": [{"url": "https://example.com", "title": "Test"}],
        }
    ]

    result = convert_to_anthropic_tool_invoke(
        tool_calls, web_search_results=web_search_results
    )

    assert len(result) == 3
    # First: server_tool_use
    assert result[0]["type"] == "server_tool_use"
    assert result[0]["id"] == "srvtoolu_01ABC123"
    # Second: web_search_tool_result
    assert result[1]["type"] == "web_search_tool_result"
    # Third: regular tool_use
    assert result[2]["type"] == "tool_use"
    assert result[2]["id"] == "toolu_01XYZ789"


def test_anthropic_messages_pt_with_server_tool_use():
    """
    Test that anthropic_messages_pt correctly reconstructs server_tool_use from provider_specific_fields.

    Fixes: https://github.com/BerriAI/litellm/issues/17737
    """
    messages = [
        {"role": "user", "content": "Search for elephant weight and add 100"},
        {
            "role": "assistant",
            "content": "Let me search for that.",
            "tool_calls": [
                {
                    "id": "srvtoolu_01ABC123",
                    "type": "function",
                    "function": {
                        "name": "web_search",
                        "arguments": '{"query": "elephant weight"}',
                    },
                },
                {
                    "id": "toolu_01XYZ789",
                    "type": "function",
                    "function": {
                        "name": "add_numbers",
                        "arguments": '{"a": 5000, "b": 100}',
                    },
                },
            ],
            "provider_specific_fields": {
                "web_search_results": [
                    {
                        "type": "web_search_tool_result",
                        "tool_use_id": "srvtoolu_01ABC123",
                        "content": [
                            {
                                "url": "https://example.com",
                                "title": "Test",
                                "snippet": "5000 kg",
                            }
                        ],
                    }
                ]
            },
        },
        {"role": "tool", "tool_call_id": "toolu_01XYZ789", "content": "5100"},
    ]

    result = anthropic_messages_pt(
        messages, model="claude-sonnet-4-5", llm_provider="anthropic"
    )

    # Find the assistant message
    assistant_msg = next(m for m in result if m["role"] == "assistant")
    content = assistant_msg["content"]

    # Should have: text, server_tool_use, web_search_tool_result, tool_use
    types = [c.get("type") for c in content]
    assert "text" in types
    assert "server_tool_use" in types
    assert "web_search_tool_result" in types
    assert "tool_use" in types

    # Verify server_tool_use
    server_tool = next(c for c in content if c.get("type") == "server_tool_use")
    assert server_tool["id"] == "srvtoolu_01ABC123"

    # Verify web_search_tool_result comes after server_tool_use
    server_idx = types.index("server_tool_use")
    web_result_idx = types.index("web_search_tool_result")
    assert web_result_idx == server_idx + 1

    # Verify regular tool_use
    tool_use = next(c for c in content if c.get("type") == "tool_use")
    assert tool_use["id"] == "toolu_01XYZ789"


def test_convert_to_anthropic_tool_invoke_with_tool_results():
    """
    Test that non-web-search *_tool_result blocks (e.g. bash_code_execution_tool_result)
    stored in provider_specific_fields["tool_results"] are paired with their server_tool_use
    block when reconstructing assistant history.

    Regression for: server tool result blocks dropped on multi-turn replay
    (bash_code_execution_tool_result, text_editor_code_execution_tool_result, etc.)
    """
    tool_calls = [
        {
            "id": "srvtoolu_01BASH",
            "type": "function",
            "function": {
                "name": "bash_code_execution",
                "arguments": '{"command": "python3 -c \\"print(2)\\""}',
            },
        }
    ]

    tool_results = [
        {
            "type": "bash_code_execution_tool_result",
            "tool_use_id": "srvtoolu_01BASH",
            "content": {
                "type": "bash_code_execution_result",
                "stdout": "2\n",
                "stderr": "",
                "return_code": 0,
                "content": [],
            },
        }
    ]

    result = convert_to_anthropic_tool_invoke(tool_calls, tool_results=tool_results)

    assert len(result) == 2
    # First: server_tool_use
    assert result[0]["type"] == "server_tool_use"
    assert result[0]["id"] == "srvtoolu_01BASH"
    assert result[0]["name"] == "bash_code_execution"
    # Second: bash_code_execution_tool_result paired correctly
    assert result[1]["type"] == "bash_code_execution_tool_result"
    assert result[1]["tool_use_id"] == "srvtoolu_01BASH"


def test_anthropic_messages_pt_raw_bash_tool_result_passthrough():
    """
    Test that raw assistant content lists containing bash_code_execution_tool_result
    blocks are passed through intact to Anthropic.

    Regression: the raw-block passthrough only handled tool_search_tool_result;
    bash_code_execution_tool_result and other *_tool_result types were silently dropped.
    """
    messages = [
        {"role": "user", "content": "What is 1+1?"},
        {
            "role": "assistant",
            "content": [
                {
                    "type": "server_tool_use",
                    "id": "srvtoolu_01BASH",
                    "name": "bash_code_execution",
                    "input": {"command": 'python3 -c "print(1+1)"'},
                },
                {
                    "type": "bash_code_execution_tool_result",
                    "tool_use_id": "srvtoolu_01BASH",
                    "content": {
                        "type": "bash_code_execution_result",
                        "stdout": "2\n",
                        "stderr": "",
                        "return_code": 0,
                        "content": [],
                    },
                },
                {"type": "text", "text": "The answer is 2."},
            ],
        },
        {"role": "user", "content": "Thanks!"},
    ]

    result = anthropic_messages_pt(
        messages, model="claude-sonnet-4-5", llm_provider="anthropic"
    )

    assistant_msg = next(m for m in result if m["role"] == "assistant")
    content = assistant_msg["content"]
    types = [c.get("type") for c in content]

    assert "server_tool_use" in types, "server_tool_use block must be preserved"
    assert (
        "bash_code_execution_tool_result" in types
    ), "bash_code_execution_tool_result block must not be dropped"
    assert "text" in types

    # Result must immediately follow its server_tool_use
    srv_idx = types.index("server_tool_use")
    result_idx = types.index("bash_code_execution_tool_result")
    assert result_idx == srv_idx + 1

    bash_result = next(
        c for c in content if c.get("type") == "bash_code_execution_tool_result"
    )
    assert bash_result["tool_use_id"] == "srvtoolu_01BASH"


def test_anthropic_messages_pt_with_bash_tool_result_in_provider_specific_fields():
    """
    Test that anthropic_messages_pt correctly reconstructs bash_code_execution_tool_result
    from provider_specific_fields["tool_results"] when replaying LiteLLM response objects.

    Regression: only web_search_results were read from provider_specific_fields;
    tool_results (bash_code_execution_tool_result, etc.) were silently lost.
    """
    messages = [
        {"role": "user", "content": "What is 1+1?"},
        {
            "role": "assistant",
            "content": None,
            "tool_calls": [
                {
                    "id": "srvtoolu_01BASH",
                    "type": "function",
                    "function": {
                        "name": "bash_code_execution",
                        "arguments": '{"command": "python3 -c \\"print(1+1)\\""}',
                    },
                }
            ],
            "provider_specific_fields": {
                "tool_results": [
                    {
                        "type": "bash_code_execution_tool_result",
                        "tool_use_id": "srvtoolu_01BASH",
                        "content": {
                            "type": "bash_code_execution_result",
                            "stdout": "2\n",
                            "stderr": "",
                            "return_code": 0,
                            "content": [],
                        },
                    }
                ]
            },
        },
        {"role": "user", "content": "Thanks!"},
    ]

    result = anthropic_messages_pt(
        messages, model="claude-sonnet-4-5", llm_provider="anthropic"
    )

    assistant_msg = next(m for m in result if m["role"] == "assistant")
    content = assistant_msg["content"]
    types = [c.get("type") for c in content]

    assert "server_tool_use" in types, "server_tool_use block must be reconstructed"
    assert (
        "bash_code_execution_tool_result" in types
    ), "bash_code_execution_tool_result must be paired from provider_specific_fields['tool_results']"

    # Result must immediately follow its server_tool_use
    srv_idx = types.index("server_tool_use")
    result_idx = types.index("bash_code_execution_tool_result")
    assert result_idx == srv_idx + 1

    srv = next(c for c in content if c.get("type") == "server_tool_use")
    assert srv["id"] == "srvtoolu_01BASH"
    bash_result = next(
        c for c in content if c.get("type") == "bash_code_execution_tool_result"
    )
    assert bash_result["tool_use_id"] == "srvtoolu_01BASH"


# ============ parse_tool_call_arguments Tests ============
# Tests for the shared utility that parses tool call JSON arguments


def test_parse_tool_call_arguments_valid_json():
    """Test that valid JSON is parsed correctly."""
    from litellm.litellm_core_utils.prompt_templates.common_utils import (
        parse_tool_call_arguments,
    )

    result = parse_tool_call_arguments('{"city": "Paris", "units": "celsius"}')
    assert result == {"city": "Paris", "units": "celsius"}


def test_parse_tool_call_arguments_empty_input():
    """Test that None/empty input returns empty dict."""
    from litellm.litellm_core_utils.prompt_templates.common_utils import (
        parse_tool_call_arguments,
    )

    assert parse_tool_call_arguments(None) == {}
    assert parse_tool_call_arguments("") == {}


def test_parse_tool_call_arguments_malformed_json():
    """Test that malformed JSON raises ValueError with context."""
    from litellm.litellm_core_utils.prompt_templates.common_utils import (
        parse_tool_call_arguments,
    )

    with pytest.raises(ValueError) as exc_info:
        parse_tool_call_arguments(
            '{"skill_name": "pptx',
            tool_name="load_skill",
            context="Anthropic tool invoke",
        )

    error_msg = str(exc_info.value)
    assert "load_skill" in error_msg
    assert "Anthropic tool invoke" in error_msg
    assert '{"skill_name": "pptx' in error_msg
    assert "Unterminated string" in error_msg


def test_convert_to_anthropic_tool_invoke_malformed_json():
    """
    Test that convert_to_anthropic_tool_invoke raises ValueError with context
    when tool arguments contain malformed JSON.

    Fixes: https://github.com/BerriAI/litellm/issues/18920
    """
    tool_calls = [
        {
            "id": "toolu_01_invalid",
            "type": "function",
            "function": {
                "name": "bad_tool",
                "arguments": '{"truncated',  # Malformed JSON
            },
        }
    ]

    with pytest.raises(ValueError) as exc_info:
        convert_to_anthropic_tool_invoke(tool_calls)

    error_msg = str(exc_info.value)
    assert "bad_tool" in error_msg
    assert '{"truncated' in error_msg


# ============ _attempt_json_repair Tests ============
# Tests for the JSON repair utility that fixes truncated tool call arguments


def test_attempt_json_repair_missing_closing_brace():
    """Repair JSON truncated with a missing closing brace (issue #22312)."""
    from litellm.litellm_core_utils.prompt_templates.common_utils import (
        _attempt_json_repair,
    )

    truncated = (
        '{"command": ["bash","-lc","find /x/repos -name \'messages.py\' -type f"]'
    )
    result = _attempt_json_repair(truncated)
    assert result is not None
    assert result["command"] == [
        "bash",
        "-lc",
        "find /x/repos -name 'messages.py' -type f",
    ]


def test_attempt_json_repair_missing_bracket_and_brace():
    """Repair JSON truncated with both missing ] and }."""
    from litellm.litellm_core_utils.prompt_templates.common_utils import (
        _attempt_json_repair,
    )

    truncated = '{"items": [1, 2, 3'
    result = _attempt_json_repair(truncated)
    assert result is not None
    assert result["items"] == [1, 2, 3]


def test_attempt_json_repair_trailing_comma():
    """Repair JSON with a trailing comma before missing close."""
    from litellm.litellm_core_utils.prompt_templates.common_utils import (
        _attempt_json_repair,
    )

    truncated = '{"a": 1, "b": 2,'
    result = _attempt_json_repair(truncated)
    assert result is not None
    assert result == {"a": 1, "b": 2}


def test_attempt_json_repair_returns_none_for_unterminated_string():
    """Cannot repair an unterminated string — returns None."""
    from litellm.litellm_core_utils.prompt_templates.common_utils import (
        _attempt_json_repair,
    )

    assert _attempt_json_repair('{"key": "incomplete value') is None


def test_attempt_json_repair_returns_none_for_valid_json():
    """Valid JSON has no unmatched brackets — returns None (no repair needed)."""
    from litellm.litellm_core_utils.prompt_templates.common_utils import (
        _attempt_json_repair,
    )

    assert _attempt_json_repair('{"key": "value"}') is None


def test_attempt_json_repair_returns_none_for_empty():
    """Empty / whitespace input returns None."""
    from litellm.litellm_core_utils.prompt_templates.common_utils import (
        _attempt_json_repair,
    )

    assert _attempt_json_repair("") is None
    assert _attempt_json_repair("   ") is None


def test_attempt_json_repair_interleaved_nesting():
    """Repair JSON with interleaved {} and [] nesting."""
    from litellm.litellm_core_utils.prompt_templates.common_utils import (
        _attempt_json_repair,
    )

    # {"a": [{"b": 2  needs }]} not ]}}
    truncated = '{"a": [{"b": 2'
    result = _attempt_json_repair(truncated)
    assert result is not None
    assert result == {"a": [{"b": 2}]}


def test_attempt_json_repair_deeply_nested():
    """Repair deeply nested truncated JSON."""
    from litellm.litellm_core_utils.prompt_templates.common_utils import (
        _attempt_json_repair,
    )

    truncated = '{"x": {"y": [1, {"z": [2, 3'
    result = _attempt_json_repair(truncated)
    assert result is not None
    assert result == {"x": {"y": [1, {"z": [2, 3]}]}}


def test_parse_tool_call_arguments_whitespace_only():
    """Whitespace-only input returns empty dict."""
    from litellm.litellm_core_utils.prompt_templates.common_utils import (
        parse_tool_call_arguments,
    )

    assert parse_tool_call_arguments("   ") == {}
    assert parse_tool_call_arguments("\n") == {}


def test_parse_tool_call_arguments_non_object_json():
    """Non-object JSON (list, string, number) is returned as-is (no wrapping)."""
    from litellm.litellm_core_utils.prompt_templates.common_utils import (
        parse_tool_call_arguments,
    )

    result = parse_tool_call_arguments("[1, 2, 3]")
    assert result == [1, 2, 3]


def test_parse_tool_call_arguments_repairs_truncated_json():
    """parse_tool_call_arguments should repair truncated JSON instead of raising."""
    from litellm.litellm_core_utils.prompt_templates.common_utils import (
        parse_tool_call_arguments,
    )

    truncated = '{"command": ["bash","-lc","find /x -type f"]'
    result = parse_tool_call_arguments(
        truncated, tool_name="shell", context="Anthropic tool invoke"
    )
    assert result == {"command": ["bash", "-lc", "find /x -type f"]}


def test_parse_tool_call_arguments_still_raises_for_unrepairable():
    """parse_tool_call_arguments raises ValueError when repair also fails."""
    from litellm.litellm_core_utils.prompt_templates.common_utils import (
        parse_tool_call_arguments,
    )

    with pytest.raises(ValueError) as exc_info:
        parse_tool_call_arguments(
            '{"key": "unterminated',
            tool_name="test_tool",
            context="test context",
        )

    error_msg = str(exc_info.value)
    assert "test_tool" in error_msg
    assert "test context" in error_msg


def test_anthropic_messages_pt_interleave_thinking_with_server_tool_calls():
    """
    Test that thinking blocks are interleaved with server tool calls (web search)
    instead of being prepended all at once.

    When Anthropic returns a response with extended thinking + multiple web searches,
    the content blocks are interleaved:
    [thinking_1, server_tool_use_1, result_1, thinking_2, server_tool_use_2, result_2]

    On round-trip through OpenAI format, thinking_blocks and tool_calls are separate
    fields. anthropic_messages_pt must reconstruct the interleaved order, otherwise
    Anthropic rejects the request because thinking block signatures are position-dependent.

    Fixes: https://github.com/BerriAI/litellm/issues/23047
    """
    messages = [
        {"role": "user", "content": "Search for news about fast.ai and answer.ai"},
        {
            "role": "assistant",
            "content": "Here is what I found.",
            "thinking_blocks": [
                {
                    "type": "thinking",
                    "thinking": "I need to search for fast.ai news.",
                    "signature": "sig_thinking_1",
                },
                {
                    "type": "thinking",
                    "thinking": "Now I should also search for answer.ai.",
                    "signature": "sig_thinking_2",
                },
            ],
            "tool_calls": [
                {
                    "id": "srvtoolu_01SEARCH1",
                    "type": "function",
                    "function": {
                        "name": "web_search",
                        "arguments": '{"query": "fast.ai news"}',
                    },
                },
                {
                    "id": "srvtoolu_01SEARCH2",
                    "type": "function",
                    "function": {
                        "name": "web_search",
                        "arguments": '{"query": "answer.ai news"}',
                    },
                },
            ],
            "provider_specific_fields": {
                "web_search_results": [
                    {
                        "type": "web_search_tool_result",
                        "tool_use_id": "srvtoolu_01SEARCH1",
                        "content": [
                            {
                                "type": "web_search_result",
                                "url": "https://fast.ai",
                                "title": "fast.ai",
                                "snippet": "fast.ai news",
                            }
                        ],
                    },
                    {
                        "type": "web_search_tool_result",
                        "tool_use_id": "srvtoolu_01SEARCH2",
                        "content": [
                            {
                                "type": "web_search_result",
                                "url": "https://answer.ai",
                                "title": "answer.ai",
                                "snippet": "answer.ai news",
                            }
                        ],
                    },
                ]
            },
        },
        {"role": "user", "content": "Now search for news about solveit"},
    ]

    result = anthropic_messages_pt(
        messages, model="claude-sonnet-4-5", llm_provider="anthropic"
    )

    # Find the assistant message
    assistant_msg = next(m for m in result if m["role"] == "assistant")
    content = assistant_msg["content"]

    # Extract types in order
    types = [c.get("type") for c in content]

    # The correct interleaved order should be:
    # thinking_1, server_tool_use_1, web_search_tool_result_1,
    # thinking_2, server_tool_use_2, web_search_tool_result_2,
    # text
    assert types == [
        "thinking",
        "server_tool_use",
        "web_search_tool_result",
        "thinking",
        "server_tool_use",
        "web_search_tool_result",
        "text",
    ], f"Expected interleaved order but got: {types}"

    # Verify thinking blocks preserved their content and signatures
    thinking_1 = content[0]
    assert thinking_1["thinking"] == "I need to search for fast.ai news."
    assert thinking_1["signature"] == "sig_thinking_1"

    thinking_2 = content[3]
    assert thinking_2["thinking"] == "Now I should also search for answer.ai."
    assert thinking_2["signature"] == "sig_thinking_2"

    # Verify server_tool_use blocks preserved their IDs
    assert content[1]["id"] == "srvtoolu_01SEARCH1"
    assert content[4]["id"] == "srvtoolu_01SEARCH2"

    # Verify web_search_tool_result blocks are paired correctly
    assert content[2]["tool_use_id"] == "srvtoolu_01SEARCH1"
    assert content[5]["tool_use_id"] == "srvtoolu_01SEARCH2"

    # Verify text block is present at the end
    assert content[6]["text"] == "Here is what I found."


def test_anthropic_messages_pt_thinking_blocks_no_server_tools_unchanged():
    """
    Test that the existing behavior is preserved when thinking blocks exist
    but there are no server tool calls (only regular tool_use).

    Thinking blocks should still be prepended first in this case.
    """
    messages = [
        {"role": "user", "content": "What is the weather?"},
        {
            "role": "assistant",
            "content": "Let me check.",
            "thinking_blocks": [
                {
                    "type": "thinking",
                    "thinking": "I should check the weather.",
                    "signature": "sig_1",
                },
            ],
            "tool_calls": [
                {
                    "id": "toolu_01REG",
                    "type": "function",
                    "function": {
                        "name": "get_weather",
                        "arguments": '{"location": "SF"}',
                    },
                },
            ],
        },
        {
            "role": "tool",
            "tool_call_id": "toolu_01REG",
            "content": "72F and sunny",
        },
    ]

    result = anthropic_messages_pt(
        messages, model="claude-sonnet-4-5", llm_provider="anthropic"
    )

    assistant_msg = next(m for m in result if m["role"] == "assistant")
    content = assistant_msg["content"]
    types = [c.get("type") for c in content]

    # Original behavior: thinking first, then text, then tool_use
    assert types == [
        "thinking",
        "text",
        "tool_use",
    ], f"Expected sequential order but got: {types}"


def test_anthropic_messages_pt_interleave_more_thinking_than_tool_groups():
    """
    Test interleaving when there are more thinking blocks than server tool groups.
    Extra thinking blocks should appear before the text block.
    """
    messages = [
        {"role": "user", "content": "Search for something"},
        {
            "role": "assistant",
            "content": "Found it.",
            "thinking_blocks": [
                {
                    "type": "thinking",
                    "thinking": "First thought",
                    "signature": "sig_1",
                },
                {
                    "type": "thinking",
                    "thinking": "Second thought",
                    "signature": "sig_2",
                },
                {
                    "type": "thinking",
                    "thinking": "Third thought after search",
                    "signature": "sig_3",
                },
            ],
            "tool_calls": [
                {
                    "id": "srvtoolu_01ONLY",
                    "type": "function",
                    "function": {
                        "name": "web_search",
                        "arguments": '{"query": "something"}',
                    },
                },
            ],
            "provider_specific_fields": {
                "web_search_results": [
                    {
                        "type": "web_search_tool_result",
                        "tool_use_id": "srvtoolu_01ONLY",
                        "content": [
                            {
                                "type": "web_search_result",
                                "url": "https://example.com",
                                "title": "Test",
                                "snippet": "result",
                            }
                        ],
                    },
                ]
            },
        },
    ]

    result = anthropic_messages_pt(
        messages, model="claude-sonnet-4-5", llm_provider="anthropic"
    )

    assistant_msg = next(m for m in result if m["role"] == "assistant")
    content = assistant_msg["content"]
    types = [c.get("type") for c in content]

    # thinking_1 paired with tool group, thinking_2 and thinking_3 before text
    assert types == [
        "thinking",  # paired with tool group
        "server_tool_use",
        "web_search_tool_result",
        "thinking",  # extra - before text
        "thinking",  # extra - before text
        "text",
    ], f"Expected order but got: {types}"


def test_anthropic_messages_pt_list_content_with_thinking_preserves_order():
    """
    Test that when assistant content is already a list containing interleaved
    thinking blocks and server tool blocks, the thinking_blocks from
    provider_specific_fields are NOT duplicated/prepended.

    This covers the gap identified by Greptile where list-content messages
    bypass INTERLEAVED MODE and fall into SEQUENTIAL MODE, which previously
    would prepend all thinking_blocks again, causing duplication and
    breaking Anthropic's position-dependent signature verification.

    Fixes: https://github.com/BerriAI/litellm/issues/23047
    """
    messages = [
        {"role": "user", "content": "Search for AI news"},
        {
            "role": "assistant",
            # Content is already a list with interleaved thinking + server tool blocks
            "content": [
                {
                    "type": "thinking",
                    "thinking": "Let me search for AI news.",
                    "signature": "sig_1",
                },
                {
                    "type": "server_tool_use",
                    "id": "srvtoolu_01SEARCH1",
                    "name": "web_search",
                    "input": {"query": "AI news"},
                },
                {
                    "type": "web_search_tool_result",
                    "tool_use_id": "srvtoolu_01SEARCH1",
                    "content": [
                        {
                            "type": "web_search_result",
                            "url": "https://example.com",
                            "title": "AI News",
                            "snippet": "Latest AI news",
                        }
                    ],
                },
                {
                    "type": "thinking",
                    "thinking": "Now let me summarize.",
                    "signature": "sig_2",
                },
                {
                    "type": "text",
                    "text": "Here is the AI news summary.",
                },
            ],
            # thinking_blocks also present in provider_specific_fields
            "thinking_blocks": [
                {
                    "type": "thinking",
                    "thinking": "Let me search for AI news.",
                    "signature": "sig_1",
                },
                {
                    "type": "thinking",
                    "thinking": "Now let me summarize.",
                    "signature": "sig_2",
                },
            ],
        },
        {"role": "user", "content": "Tell me more"},
    ]

    result = anthropic_messages_pt(
        messages, model="claude-sonnet-4-5", llm_provider="anthropic"
    )

    assistant_msg = next(m for m in result if m["role"] == "assistant")
    content = assistant_msg["content"]
    types = [c.get("type") for c in content]

    # The list content already has the correct interleaved order.
    # thinking_blocks should NOT be prepended again (which would cause
    # duplication and break signature verification).
    assert types == [
        "thinking",
        "server_tool_use",
        "web_search_tool_result",
        "thinking",
        "text",
    ], f"Expected preserved list order without duplicate thinking blocks, but got: {types}"

    # Verify no duplicate thinking blocks
    thinking_count = sum(1 for t in types if t == "thinking")
    assert (
        thinking_count == 2
    ), f"Expected 2 thinking blocks, got {thinking_count} (duplication detected)"

    # Verify signatures preserved in correct positions
    assert content[0]["signature"] == "sig_1"
    assert content[3]["signature"] == "sig_2"
