import vertexai.generative_models as genai
from agentdojo.agent_pipeline.llms.google_llm import (
    _merge_tool_result_messages,
)


def test_merge_tool_result_messages():
    messages = [
        genai.Content(
            role="user",
            parts=[genai.Part.from_text("Hi")],
        ),
        genai.Content(
            parts=[
                genai.Part.from_function_response(
                    name="tool1",
                    response={"content": "output1"},
                )
            ],
        ),
        genai.Content(
            parts=[
                genai.Part.from_function_response(
                    name="tool2",
                    response={"content": "output2"},
                )
            ],
        ),
        genai.Content(
            role="user",
            parts=[genai.Part.from_text("Bye")],
        ),
        genai.Content(
            parts=[
                genai.Part.from_function_response(
                    name="tool3",
                    response={"content": "output3"},
                )
            ],
        ),
        genai.Content(
            parts=[
                genai.Part.from_function_response(
                    name="tool4",
                    response={"content": "output4"},
                )
            ],
        ),
    ]

    expected_messages = [
        genai.Content(
            role="user",
            parts=[genai.Part.from_text("Hi")],
        ),
        genai.Content(
            parts=[
                genai.Part.from_function_response(
                    name="tool1",
                    response={"content": "output1"},
                ),
                genai.Part.from_function_response(
                    name="tool2",
                    response={"content": "output2"},
                ),
            ],
        ),
        genai.Content(
            role="user",
            parts=[genai.Part.from_text("Bye")],
        ),
        genai.Content(
            parts=[
                genai.Part.from_function_response(
                    name="tool3",
                    response={"content": "output3"},
                ),
                genai.Part.from_function_response(
                    name="tool4",
                    response={"content": "output4"},
                ),
            ],
        ),
    ]
    merged_messages = _merge_tool_result_messages(messages)

    assert [message.to_dict() for message in merged_messages] == [message.to_dict() for message in expected_messages]
