    response_mock.usage.completion_tokens = output_tokens
    input_tokens_details_mock = MagicMock()
    input_tokens_details_mock.cached_tokens = 0
    response_mock.usage.input_tokens_details = input_tokens_details_mock

    return response_mock

# --- Test MessageBuilders ---

def test_openai_response_api_message_builder_text():
    builder = OpenAIResponseAPIMessageBuilder.user()
    builder.add_text("Hello, world!")
    messages = builder.prepare_message()
    assert len(messages) == 1
    assert messages[0]["role"] == "user"
    assert messages[0]["content"] == [{"type": "input_text", "text": "Hello, world!"}]

def test_openai_response_api_message_builder_image():
    builder = OpenAIResponseAPIMessageBuilder.user()
    builder.add_image("data:image/png;base64,SIMPLEBASE64STRING")
    messages = builder.prepare_message()
    assert len(messages) == 1
    assert messages[0]["role"] == "user"
    assert messages[0]["content"] == [
        {"type": "input_image", "image_url": "data:image/png;base64,SIMPLEBASE64STRING"}
    ]

def test_anthropic_api_message_builder_text():
    builder = AnthropicAPIMessageBuilder.user()
    builder.add_text("Hello, Anthropic!")
    messages = builder.prepare_message()
    assert len(messages) == 1
    assert messages[0]["role"] == "user"
    assert messages[0]["content"] == [{"type": "text", "text": "Hello, Anthropic!"}]

def test_anthropic_api_message_builder_image():
    builder = AnthropicAPIMessageBuilder.user()
    builder.add_image("data:image/png;base64,ANTHROPICBASE64")
    messages = builder.prepare_message()
    assert len(messages) == 1
    assert messages[0]["role"] == "user"
    assert len(messages[0]["content"]) == 1
    image_content = messages[0]["content"][0]
    assert image_content["type"] == "image"


def test_openai_response_api_message_builder_text():
    builder = OpenAIResponseAPIMessageBuilder.user()
    builder.add_text("Hello, world!")
    messages = builder.prepare_message()
    assert len(messages) == 1
    assert messages[0]["role"] == "user"
    assert messages[0]["content"] == [{"type": "input_text", "text": "Hello, world!"}]

def test_openai_response_api_message_builder_image():
    builder = OpenAIResponseAPIMessageBuilder.user()
    builder.add_image("data:image/png;base64,SIMPLEBASE64STRING")
    messages = builder.prepare_message()
    assert len(messages) == 1
    assert messages[0]["role"] == "user"
    assert messages[0]["content"] == [
        {"type": "input_image", "image_url": "data:image/png;base64,SIMPLEBASE64STRING"}
    ]

def test_anthropic_api_message_builder_text():
    builder = AnthropicAPIMessageBuilder.user()
    builder.add_text("Hello, Anthropic!")
    messages = builder.prepare_message()
    assert len(messages) == 1
    assert messages[0]["role"] == "user"
    assert messages[0]["content"] == [{"type": "text", "text": "Hello, Anthropic!"}]

def test_anthropic_api_message_builder_image():
    builder = AnthropicAPIMessageBuilder.user()
    builder.add_image("data:image/png;base64,ANTHROPICBASE64")
    messages = builder.prepare_message()
    assert len(messages) == 1
    assert messages[0]["role"] == "user"
    assert len(messages[0]["content"]) == 1
    image_content = messages[0]["content"][0]
    assert image_content["type"] == "image"
    assert image_content["source"]["type"] == "base64"
    assert image_content["source"]["media_type"] == "image/png"
    assert image_content["source"]["data"] == "ANTHROPICBASE64"  # Base64 prefix should be stripped

def test_openai_chat_completion_api_message_builder_text():
    builder = OpenAIChatCompletionAPIMessageBuilder.user()
    builder.add_text("Hello, ChatCompletion!")
    messages = builder.prepare_message()


def test_openai_response_api_message_builder_image():
    builder = OpenAIResponseAPIMessageBuilder.user()
    builder.add_image("data:image/png;base64,SIMPLEBASE64STRING")
    messages = builder.prepare_message()
    assert len(messages) == 1
    assert messages[0]["role"] == "user"
    assert messages[0]["content"] == [
        {"type": "input_image", "image_url": "data:image/png;base64,SIMPLEBASE64STRING"}
    ]

def test_anthropic_api_message_builder_text():
    builder = AnthropicAPIMessageBuilder.user()
    builder.add_text("Hello, Anthropic!")
    messages = builder.prepare_message()
    assert len(messages) == 1
    assert messages[0]["role"] == "user"
    assert messages[0]["content"] == [{"type": "text", "text": "Hello, Anthropic!"}]

def test_anthropic_api_message_builder_image():
    builder = AnthropicAPIMessageBuilder.user()
    builder.add_image("data:image/png;base64,ANTHROPICBASE64")
    messages = builder.prepare_message()
    assert len(messages) == 1
    assert messages[0]["role"] == "user"
    assert len(messages[0]["content"]) == 1
    image_content = messages[0]["content"][0]
    assert image_content["type"] == "image"
    assert image_content["source"]["type"] == "base64"
    assert image_content["source"]["media_type"] == "image/png"
    assert image_content["source"]["data"] == "ANTHROPICBASE64"  # Base64 prefix should be stripped

def test_openai_chat_completion_api_message_builder_text():
    builder = OpenAIChatCompletionAPIMessageBuilder.user()
    builder.add_text("Hello, ChatCompletion!")
    messages = builder.prepare_message()

    assert len(messages) == 1
    assert messages[0]["role"] == "user"
    assert messages[0]["content"] == [{"type": "text", "text": "Hello, ChatCompletion!"}]

def test_openai_chat_completion_api_message_builder_image():
    builder = OpenAIChatCompletionAPIMessageBuilder.user()
    builder.add_image("data:image/jpeg;base64,CHATCOMPLETIONBASE64")
    messages = builder.prepare_message()

    assert len(messages) == 1
    assert messages[0]["role"] == "user"
    assert messages[0]["content"] == [


    messages = builder.prepare_message()
    assert len(messages) == 1
    assert messages[0]["role"] == "user"
    assert messages[0]["content"] == [{"type": "text", "text": "Hello, Anthropic!"}]

def test_anthropic_api_message_builder_image():
    builder = AnthropicAPIMessageBuilder.user()
    builder.add_image("data:image/png;base64,ANTHROPICBASE64")
    messages = builder.prepare_message()
    assert len(messages) == 1
    assert messages[0]["role"] == "user"
    assert len(messages[0]["content"]) == 1
    image_content = messages[0]["content"][0]
    assert image_content["type"] == "image"
    assert image_content["source"]["type"] == "base64"
    assert image_content["source"]["media_type"] == "image/png"
    assert image_content["source"]["data"] == "ANTHROPICBASE64"  # Base64 prefix should be stripped

def test_openai_chat_completion_api_message_builder_text():
    builder = OpenAIChatCompletionAPIMessageBuilder.user()
    builder.add_text("Hello, ChatCompletion!")
    messages = builder.prepare_message()

    assert len(messages) == 1
    assert messages[0]["role"] == "user"
    assert messages[0]["content"] == [{"type": "text", "text": "Hello, ChatCompletion!"}]

def test_openai_chat_completion_api_message_builder_image():
    builder = OpenAIChatCompletionAPIMessageBuilder.user()
    builder.add_image("data:image/jpeg;base64,CHATCOMPLETIONBASE64")
    messages = builder.prepare_message()

    assert len(messages) == 1
    assert messages[0]["role"] == "user"
    assert messages[0]["content"] == [
        {"type": "image_url", "image_url": {"url": "data:image/jpeg;base64,CHATCOMPLETIONBASE64"}}
    ]

def test_openai_chat_completion_model_parse_and_cost():
    args = OpenAIChatModelArgs(model_name="gpt-3.5-turbo")
    with patch("agentlab.llm.response_api.OpenAI") as mock_openai_class:
        mock_client = MagicMock()
        mock_openai_class.return_value = mock_client
        model = args.make_model()


    assert len(messages) == 1
    assert messages[0]["role"] == "user"
    assert len(messages[0]["content"]) == 1
    image_content = messages[0]["content"][0]
    assert image_content["type"] == "image"
    assert image_content["source"]["type"] == "base64"
    assert image_content["source"]["media_type"] == "image/png"
    assert image_content["source"]["data"] == "ANTHROPICBASE64"  # Base64 prefix should be stripped

def test_openai_chat_completion_api_message_builder_text():
    builder = OpenAIChatCompletionAPIMessageBuilder.user()
    builder.add_text("Hello, ChatCompletion!")
    messages = builder.prepare_message()

    assert len(messages) == 1
    assert messages[0]["role"] == "user"
    assert messages[0]["content"] == [{"type": "text", "text": "Hello, ChatCompletion!"}]

def test_openai_chat_completion_api_message_builder_image():
    builder = OpenAIChatCompletionAPIMessageBuilder.user()
    builder.add_image("data:image/jpeg;base64,CHATCOMPLETIONBASE64")
    messages = builder.prepare_message()

    assert len(messages) == 1
    assert messages[0]["role"] == "user"
    assert messages[0]["content"] == [
        {"type": "image_url", "image_url": {"url": "data:image/jpeg;base64,CHATCOMPLETIONBASE64"}}
    ]

def test_openai_chat_completion_model_parse_and_cost():
    args = OpenAIChatModelArgs(model_name="gpt-3.5-turbo")
    with patch("agentlab.llm.response_api.OpenAI") as mock_openai_class:
        mock_client = MagicMock()
        mock_openai_class.return_value = mock_client
        model = args.make_model()

    mock_response = create_mock_openai_chat_completion(
        content="This is a test thought.",
        tool_calls=[
            {
                "id": "call_123",
                "type": "function",
                "function": {"name": "get_weather", "arguments": '{"location": "Paris"}'},
            }
        ],
        prompt_tokens=50,
        completion_tokens=30,


    builder.add_text("Hello, ChatCompletion!")
    messages = builder.prepare_message()

    assert len(messages) == 1
    assert messages[0]["role"] == "user"
    assert messages[0]["content"] == [{"type": "text", "text": "Hello, ChatCompletion!"}]

def test_openai_chat_completion_api_message_builder_image():
    builder = OpenAIChatCompletionAPIMessageBuilder.user()
    builder.add_image("data:image/jpeg;base64,CHATCOMPLETIONBASE64")
    messages = builder.prepare_message()

    assert len(messages) == 1
    assert messages[0]["role"] == "user"
    assert messages[0]["content"] == [
        {"type": "image_url", "image_url": {"url": "data:image/jpeg;base64,CHATCOMPLETIONBASE64"}}
    ]

def test_openai_chat_completion_model_parse_and_cost():
    args = OpenAIChatModelArgs(model_name="gpt-3.5-turbo")
    with patch("agentlab.llm.response_api.OpenAI") as mock_openai_class:
        mock_client = MagicMock()
        mock_openai_class.return_value = mock_client
        model = args.make_model()

    mock_response = create_mock_openai_chat_completion(
        content="This is a test thought.",
        tool_calls=[
            {
                "id": "call_123",
                "type": "function",
                "function": {"name": "get_weather", "arguments": '{"location": "Paris"}'},
            }
        ],
        prompt_tokens=50,
        completion_tokens=30,
    )

    with patch.object(
        model.client.chat.completions, "create", return_value=mock_response
    ) as mock_create:
        with tracking.set_tracker() as global_tracker:
            messages = [
                OpenAIChatCompletionAPIMessageBuilder.user().add_text(
                    "What's the weather in Paris?"
                )
            ]
            payload = APIPayload(messages=messages)
            parsed_output = model(payload)

    mock_create.assert_called_once()
    assert parsed_output.raw_response.choices[0].message.content == "This is a test thought."
    assert parsed_output.action == """get_weather(location='Paris')"""
    assert parsed_output.raw_response.choices[0].message.tool_calls[0].id == "call_123"
    # Check cost tracking (token counts)
    assert global_tracker.stats["input_tokens"] == 50
    assert global_tracker.stats["output_tokens"] == 30
    assert global_tracker.stats["cost"] > 0

def test_claude_response_model_parse_and_cost():
    args = ClaudeResponseModelArgs(model_name="claude-3-haiku-20240307")
    model = args.make_model()

    mock_anthropic_api_response = create_mock_anthropic_response(
        text_content="Thinking about the request.",
        tool_use={"id": "tool_abc", "name": "search_web", "input": {"query": "latest news"}},
        input_tokens=40,
        output_tokens=20,
    )

    with patch.object(
        model.client.messages, "create", return_value=mock_anthropic_api_response
    ) as mock_create:
        with tracking.set_tracker() as global_tracker:
            messages = [AnthropicAPIMessageBuilder.user().add_text("Search for latest news")]
            payload = APIPayload(messages=messages)
            parsed_output = model(payload)


    ) as mock_create:
        with tracking.set_tracker() as global_tracker:
            messages = [
                OpenAIChatCompletionAPIMessageBuilder.user().add_text(
                    "What's the weather in Paris?"
                )
            ]
            payload = APIPayload(messages=messages)
            parsed_output = model(payload)

    mock_create.assert_called_once()
    assert parsed_output.raw_response.choices[0].message.content == "This is a test thought."
    assert parsed_output.action == """get_weather(location='Paris')"""
    assert parsed_output.raw_response.choices[0].message.tool_calls[0].id == "call_123"
    # Check cost tracking (token counts)
    assert global_tracker.stats["input_tokens"] == 50
    assert global_tracker.stats["output_tokens"] == 30
    assert global_tracker.stats["cost"] > 0

def test_claude_response_model_parse_and_cost():
    args = ClaudeResponseModelArgs(model_name="claude-3-haiku-20240307")
    model = args.make_model()

    mock_anthropic_api_response = create_mock_anthropic_response(
        text_content="Thinking about the request.",
        tool_use={"id": "tool_abc", "name": "search_web", "input": {"query": "latest news"}},
        input_tokens=40,
        output_tokens=20,
    )

    with patch.object(
        model.client.messages, "create", return_value=mock_anthropic_api_response
    ) as mock_create:
        with tracking.set_tracker() as global_tracker:
            messages = [AnthropicAPIMessageBuilder.user().add_text("Search for latest news")]
            payload = APIPayload(messages=messages)
            parsed_output = model(payload)

    mock_create.assert_called_once()
    fn_call = next(iter(parsed_output.tool_calls))

    assert "Thinking about the request." in parsed_output.think
    assert parsed_output.action == """search_web(query='latest news')"""
    assert fn_call.name == "search_web"
    assert global_tracker.stats["input_tokens"] == 40
    assert global_tracker.stats["output_tokens"] == 20

def test_openai_response_model_parse_and_cost():
    args = OpenAIResponseModelArgs(model_name="gpt-4.1")

    mock_function_call_output = {
        "type": "function_call",
        "name": "get_current_weather",
        "arguments": '{"location": "Boston, MA", "unit": "celsius"}',
        "call_id": "call_abc123",
    }

    mock_api_resp = create_mock_openai_responses_api_response(
        outputs=[mock_function_call_output],
        input_tokens=70,
        output_tokens=40,
    )

    with patch("agentlab.llm.response_api.OpenAI") as mock_openai_class:
        mock_client = MagicMock()


    )

    with patch.object(
        model.client.messages, "create", return_value=mock_anthropic_api_response
    ) as mock_create:
        with tracking.set_tracker() as global_tracker:
            messages = [AnthropicAPIMessageBuilder.user().add_text("Search for latest news")]
            payload = APIPayload(messages=messages)
            parsed_output = model(payload)

    mock_create.assert_called_once()
    fn_call = next(iter(parsed_output.tool_calls))

    assert "Thinking about the request." in parsed_output.think
    assert parsed_output.action == """search_web(query='latest news')"""
    assert fn_call.name == "search_web"
    assert global_tracker.stats["input_tokens"] == 40
    assert global_tracker.stats["output_tokens"] == 20

def test_openai_response_model_parse_and_cost():
    args = OpenAIResponseModelArgs(model_name="gpt-4.1")

    mock_function_call_output = {
        "type": "function_call",
        "name": "get_current_weather",
        "arguments": '{"location": "Boston, MA", "unit": "celsius"}',
        "call_id": "call_abc123",
    }

    mock_api_resp = create_mock_openai_responses_api_response(
        outputs=[mock_function_call_output],
        input_tokens=70,
        output_tokens=40,
    )

    with patch("agentlab.llm.response_api.OpenAI") as mock_openai_class:
        mock_client = MagicMock()
        mock_openai_class.return_value = mock_client
        model = args.make_model()

    with patch.object(
        model.client.responses, "create", return_value=mock_api_resp
    ) as mock_create_method:
        with tracking.set_tracker() as global_tracker:
            messages = [
                OpenAIResponseAPIMessageBuilder.user().add_text("What's the weather in Boston?")
            ]
            payload = APIPayload(messages=messages)
            parsed_output = model(payload)

    mock_create_method.assert_called_once()
    fn_calls = [
        content
        for content in parsed_output.tool_calls.raw_calls.output
        if content.type == "function_call"
    ]
    assert parsed_output.action == "get_current_weather(location='Boston, MA', unit='celsius')"
    assert fn_calls[0].call_id == "call_abc123"
    assert parsed_output.raw_response == mock_api_resp
    assert global_tracker.stats["input_tokens"] == 70
    assert global_tracker.stats["output_tokens"] == 40

# --- Test Response Models (Pricy - require API keys and actual calls) ---

@pytest.mark.pricy
@pytest.mark.skipif(not os.getenv("OPENAI_API_KEY"), reason="OPENAI_API_KEY not set")
def test_openai_chat_completion_model_pricy_call():
    """Tests OpenAIChatCompletionModel with a real API call."""
    args = OpenAIChatModelArgs(
        model_name="gpt-4.1",
        temperature=1e-5,
        max_new_tokens=100,
    )

    tools = chat_api_tools
    model = args.make_model()

    with tracking.set_tracker() as global_tracker:
        messages = [


            parsed_output = model(payload)

    mock_create_method.assert_called_once()
    fn_calls = [
        content
        for content in parsed_output.tool_calls.raw_calls.output
        if content.type == "function_call"
    ]
    assert parsed_output.action == "get_current_weather(location='Boston, MA', unit='celsius')"
    assert fn_calls[0].call_id == "call_abc123"
    assert parsed_output.raw_response == mock_api_resp
    assert global_tracker.stats["input_tokens"] == 70
    assert global_tracker.stats["output_tokens"] == 40

# --- Test Response Models (Pricy - require API keys and actual calls) ---

@pytest.mark.pricy
@pytest.mark.skipif(not os.getenv("OPENAI_API_KEY"), reason="OPENAI_API_KEY not set")
def test_openai_chat_completion_model_pricy_call():
    """Tests OpenAIChatCompletionModel with a real API call."""
    args = OpenAIChatModelArgs(
        model_name="gpt-4.1",
        temperature=1e-5,
        max_new_tokens=100,
    )

    tools = chat_api_tools
    model = args.make_model()

    with tracking.set_tracker() as global_tracker:
        messages = [
            OpenAIChatCompletionAPIMessageBuilder.user().add_text("What is the weather in Paris?")
        ]
        payload = APIPayload(messages=messages, tools=tools, tool_choice="required")
        parsed_output = model(payload)

    assert parsed_output.raw_response is not None
    assert (
        parsed_output.action == "get_weather(location='Paris')"
    ), f""" Expected get_weather(location='Paris') but got {parsed_output.action}"""
    assert global_tracker.stats["input_tokens"] > 0
    assert global_tracker.stats["output_tokens"] > 0
    assert global_tracker.stats["cost"] > 0

@pytest.mark.pricy
@pytest.mark.skipif(not os.getenv("ANTHROPIC_API_KEY"), reason="ANTHROPIC_API_KEY not set")
def test_claude_response_model_pricy_call():
    """Tests ClaudeResponseModel with a real API call."""

    args = ClaudeResponseModelArgs(
        model_name="claude-3-haiku-20240307",
        temperature=1e-5,
        max_new_tokens=100,
    )
    tools = anthropic_tools
    model = args.make_model()

    with tracking.set_tracker() as global_tracker:
        messages = [AnthropicAPIMessageBuilder.user().add_text("What is the weather in Paris?")]
        payload = APIPayload(messages=messages, tools=tools)
        parsed_output = model(payload)


    model = args.make_model()

    with tracking.set_tracker() as global_tracker:
        messages = [
            OpenAIChatCompletionAPIMessageBuilder.user().add_text("What is the weather in Paris?")
        ]
        payload = APIPayload(messages=messages, tools=tools, tool_choice="required")
        parsed_output = model(payload)

    assert parsed_output.raw_response is not None
    assert (
        parsed_output.action == "get_weather(location='Paris')"
    ), f""" Expected get_weather(location='Paris') but got {parsed_output.action}"""
    assert global_tracker.stats["input_tokens"] > 0
    assert global_tracker.stats["output_tokens"] > 0
    assert global_tracker.stats["cost"] > 0

@pytest.mark.pricy
@pytest.mark.skipif(not os.getenv("ANTHROPIC_API_KEY"), reason="ANTHROPIC_API_KEY not set")
def test_claude_response_model_pricy_call():
    """Tests ClaudeResponseModel with a real API call."""

    args = ClaudeResponseModelArgs(
        model_name="claude-3-haiku-20240307",
        temperature=1e-5,
        max_new_tokens=100,
    )
    tools = anthropic_tools
    model = args.make_model()

    with tracking.set_tracker() as global_tracker:
        messages = [AnthropicAPIMessageBuilder.user().add_text("What is the weather in Paris?")]
        payload = APIPayload(messages=messages, tools=tools)
        parsed_output = model(payload)

    assert parsed_output.raw_response is not None
    assert (
        parsed_output.action == "get_weather(location='Paris')"
    ), f"""Expected get_weather('Paris') but got {parsed_output.action}"""
    assert global_tracker.stats["input_tokens"] > 0
    assert global_tracker.stats["output_tokens"] > 0
    assert global_tracker.stats["cost"] > 0

@pytest.mark.pricy
@pytest.mark.skipif(not os.getenv("OPENAI_API_KEY"), reason="OPENAI_API_KEY not set")
def test_openai_response_model_pricy_call():
    """
    Tests OpenAIResponseModel output parsing and cost tracking with both
    function_call and reasoning outputs.
    """
    args = OpenAIResponseModelArgs(model_name="gpt-4.1", temperature=1e-5, max_new_tokens=100)

    tools = responses_api_tools
    model = args.make_model()

    with tracking.set_tracker() as global_tracker:
        messages = [
            OpenAIResponseAPIMessageBuilder.user().add_text("What is the weather in Paris?")
        ]
        payload = APIPayload(messages=messages, tools=tools)
        parsed_output = model(payload)


    )
    tools = anthropic_tools
    model = args.make_model()

    with tracking.set_tracker() as global_tracker:
        messages = [AnthropicAPIMessageBuilder.user().add_text("What is the weather in Paris?")]
        payload = APIPayload(messages=messages, tools=tools)
        parsed_output = model(payload)

    assert parsed_output.raw_response is not None
    assert (
        parsed_output.action == "get_weather(location='Paris')"
    ), f"""Expected get_weather('Paris') but got {parsed_output.action}"""
    assert global_tracker.stats["input_tokens"] > 0
    assert global_tracker.stats["output_tokens"] > 0
    assert global_tracker.stats["cost"] > 0

@pytest.mark.pricy
@pytest.mark.skipif(not os.getenv("OPENAI_API_KEY"), reason="OPENAI_API_KEY not set")
def test_openai_response_model_pricy_call():
    """
    Tests OpenAIResponseModel output parsing and cost tracking with both
    function_call and reasoning outputs.
    """
    args = OpenAIResponseModelArgs(model_name="gpt-4.1", temperature=1e-5, max_new_tokens=100)

    tools = responses_api_tools
    model = args.make_model()

    with tracking.set_tracker() as global_tracker:
        messages = [
            OpenAIResponseAPIMessageBuilder.user().add_text("What is the weather in Paris?")
        ]
        payload = APIPayload(messages=messages, tools=tools)
        parsed_output = model(payload)

    assert parsed_output.raw_response is not None
    assert (
        parsed_output.action == """get_weather(location='Paris', unit='celsius')"""
    ), f""" Expected get_weather(location='Paris', unit='celsius') but got {parsed_output.action}"""
    assert global_tracker.stats["input_tokens"] > 0
    assert global_tracker.stats["output_tokens"] > 0
    assert global_tracker.stats["cost"] > 0

@pytest.mark.pricy
@pytest.mark.skipif(not os.getenv("OPENAI_API_KEY"), reason="OPENAI_API_KEY not set")
def test_openai_response_model_with_multiple_messages_and_cost_tracking():
    """
    Test OpenAIResponseModel's output parsing and cost tracking
    with a tool-using assistant and follow-up interaction.
    """
    args = OpenAIResponseModelArgs(model_name="gpt-4.1", temperature=1e-5, max_new_tokens=100)

    tools = responses_api_tools
    model = args.make_model()
    builder = args.get_message_builder()

    messages = [builder.user().add_text("What is the weather in Paris?")]

    with tracking.set_tracker() as tracker:
        payload = APIPayload(messages=messages, tools=tools, tool_choice="required")
        parsed = model(payload)


    model = args.make_model()

    with tracking.set_tracker() as global_tracker:
        messages = [
            OpenAIResponseAPIMessageBuilder.user().add_text("What is the weather in Paris?")
        ]
        payload = APIPayload(messages=messages, tools=tools)
        parsed_output = model(payload)

    assert parsed_output.raw_response is not None
    assert (
        parsed_output.action == """get_weather(location='Paris', unit='celsius')"""
    ), f""" Expected get_weather(location='Paris', unit='celsius') but got {parsed_output.action}"""
    assert global_tracker.stats["input_tokens"] > 0
    assert global_tracker.stats["output_tokens"] > 0
    assert global_tracker.stats["cost"] > 0

@pytest.mark.pricy
@pytest.mark.skipif(not os.getenv("OPENAI_API_KEY"), reason="OPENAI_API_KEY not set")
def test_openai_response_model_with_multiple_messages_and_cost_tracking():
    """
    Test OpenAIResponseModel's output parsing and cost tracking
    with a tool-using assistant and follow-up interaction.
    """
    args = OpenAIResponseModelArgs(model_name="gpt-4.1", temperature=1e-5, max_new_tokens=100)

    tools = responses_api_tools
    model = args.make_model()
    builder = args.get_message_builder()

    messages = [builder.user().add_text("What is the weather in Paris?")]

    with tracking.set_tracker() as tracker:
        payload = APIPayload(messages=messages, tools=tools, tool_choice="required")
        parsed = model(payload)
        prev_input = tracker.stats["input_tokens"]
        prev_output = tracker.stats["output_tokens"]
        prev_cost = tracker.stats["cost"]

        assert parsed.tool_calls, "Expected tool calls in the response"
        # Set tool responses
        for tool_call in parsed.tool_calls:
            tool_call.response_text("Its sunny! 25°C")
        # Simulate tool execution and user follow-up
        messages += [
            builder.add_responded_tool_calls(parsed.tool_calls),
            builder.user().add_text("What is the weather in Delhi?"),
        ]

        payload = APIPayload(messages=messages, tools=tools, tool_choice="required")
        parsed = model(payload)

        delta_input = tracker.stats["input_tokens"] - prev_input
        delta_output = tracker.stats["output_tokens"] - prev_output
        delta_cost = tracker.stats["cost"] - prev_cost

    assert prev_input > 0
    assert prev_output > 0
    assert prev_cost > 0
    assert parsed.raw_response is not None
    assert (
        parsed.action == """get_weather(location='Delhi', unit='celsius')"""
    ), f"Unexpected action: {parsed.action}"
    assert delta_input > 0
    assert delta_output > 0
    assert delta_cost > 0
    assert tracker.stats["input_tokens"] == prev_input + delta_input
    assert tracker.stats["output_tokens"] == prev_output + delta_output
    assert tracker.stats["cost"] == pytest.approx(prev_cost + delta_cost)

@pytest.mark.pricy
@pytest.mark.skipif(not os.getenv("OPENAI_API_KEY"), reason="OPENAI_API_KEY not set")
def test_openai_chat_completion_model_with_multiple_messages_and_cost_tracking():
    """
    Test OpenAIResponseModel's output parsing and cost tracking
    with a tool-using assistant and follow-up interaction.
    """
    args = OpenAIChatModelArgs(model_name="gpt-4.1", temperature=1e-5, max_new_tokens=100)

    tools = [
        {
            "type": "function",
            "name": "get_weather",
            "description": "Get the current weather in a given location.",
            "parameters": {
                "type": "object",
                "properties": {
                    "location": {


        delta_output = tracker.stats["output_tokens"] - prev_output
        delta_cost = tracker.stats["cost"] - prev_cost

    assert prev_input > 0
    assert prev_output > 0
    assert prev_cost > 0
    assert parsed.raw_response is not None
    assert (
        parsed.action == """get_weather(location='Delhi', unit='celsius')"""
    ), f"Unexpected action: {parsed.action}"
    assert delta_input > 0
    assert delta_output > 0
    assert delta_cost > 0
    assert tracker.stats["input_tokens"] == prev_input + delta_input
    assert tracker.stats["output_tokens"] == prev_output + delta_output
    assert tracker.stats["cost"] == pytest.approx(prev_cost + delta_cost)

@pytest.mark.pricy
@pytest.mark.skipif(not os.getenv("OPENAI_API_KEY"), reason="OPENAI_API_KEY not set")
def test_openai_chat_completion_model_with_multiple_messages_and_cost_tracking():
    """
    Test OpenAIResponseModel's output parsing and cost tracking
    with a tool-using assistant and follow-up interaction.
    """
    args = OpenAIChatModelArgs(model_name="gpt-4.1", temperature=1e-5, max_new_tokens=100)

    tools = [
        {
            "type": "function",
            "name": "get_weather",
            "description": "Get the current weather in a given location.",
            "parameters": {
                "type": "object",
                "properties": {
                    "location": {
                        "type": "string",
                        "description": "The location to get the weather for.",
                    },
                    "unit": {
                        "type": "string",
                        "enum": ["celsius", "fahrenheit"],
                        "description": "The unit of temperature.",
                    },
                },
                "required": ["location"],
            },
        }
    ]

    model = args.make_model()
    builder = args.get_message_builder()

    messages = [builder.user().add_text("What is the weather in Paris?")]

    with tracking.set_tracker() as tracker:
        payload = APIPayload(messages=messages, tools=tools, tool_choice="required")
        parsed = model(payload)
        prev_input = tracker.stats["input_tokens"]
        prev_output = tracker.stats["output_tokens"]
        prev_cost = tracker.stats["cost"]

        for tool_call in parsed.tool_calls:
            tool_call.response_text("Its sunny! 25°C")
        # Simulate tool execution and user follow-up
        messages += [
            builder.add_responded_tool_calls(parsed.tool_calls),
            builder.user().add_text("What is the weather in Delhi?"),
        ]
        # Set tool responses

        payload = APIPayload(messages=messages, tools=tools, tool_choice="required")
        parsed = model(payload)

        delta_input = tracker.stats["input_tokens"] - prev_input
        delta_output = tracker.stats["output_tokens"] - prev_output
        delta_cost = tracker.stats["cost"] - prev_cost

    assert prev_input > 0
    assert prev_output > 0
    assert prev_cost > 0
    assert parsed.raw_response is not None
    assert (
        parsed.action == """get_weather(location='Delhi')"""
    ), f"Unexpected action: {parsed.action}"
    assert delta_input > 0
    assert delta_output > 0
    assert delta_cost > 0
    assert tracker.stats["input_tokens"] == prev_input + delta_input
    assert tracker.stats["output_tokens"] == prev_output + delta_output
    assert tracker.stats["cost"] == pytest.approx(prev_cost + delta_cost)

@pytest.mark.pricy
@pytest.mark.skipif(not os.getenv("ANTHROPIC_API_KEY"), reason="ANTHROPIC_API_KEY not set")
def test_claude_model_with_multiple_messages_pricy_call():
    model_factory = ClaudeResponseModelArgs(
        model_name="claude-3-haiku-20240307", temperature=1e-5, max_new_tokens=100
    )
    tools = [
        {
            "name": "get_weather",
            "description": "Get the current weather in a given location.",
            "input_schema": {
                "type": "object",
                "properties": {
                    "location": {
                        "type": "string",
                        "description": "The location to get the weather for.",
                    },
                    "unit": {


        delta_output = tracker.stats["output_tokens"] - prev_output
        delta_cost = tracker.stats["cost"] - prev_cost

    assert prev_input > 0
    assert prev_output > 0
    assert prev_cost > 0
    assert parsed.raw_response is not None
    assert (
        parsed.action == """get_weather(location='Delhi')"""
    ), f"Unexpected action: {parsed.action}"
    assert delta_input > 0
    assert delta_output > 0
    assert delta_cost > 0
    assert tracker.stats["input_tokens"] == prev_input + delta_input
    assert tracker.stats["output_tokens"] == prev_output + delta_output
    assert tracker.stats["cost"] == pytest.approx(prev_cost + delta_cost)

@pytest.mark.pricy
@pytest.mark.skipif(not os.getenv("ANTHROPIC_API_KEY"), reason="ANTHROPIC_API_KEY not set")
def test_claude_model_with_multiple_messages_pricy_call():
    model_factory = ClaudeResponseModelArgs(
        model_name="claude-3-haiku-20240307", temperature=1e-5, max_new_tokens=100
    )
    tools = [
        {
            "name": "get_weather",
            "description": "Get the current weather in a given location.",
            "input_schema": {
                "type": "object",
                "properties": {
                    "location": {
                        "type": "string",
                        "description": "The location to get the weather for.",
                    },
                    "unit": {
                        "type": "string",
                        "enum": ["celsius", "fahrenheit"],
                        "description": "The unit of temperature.",
                    },
                },
                "required": ["location"],
            },
        }
    ]
    model = model_factory.make_model()
    msg_builder = model_factory.get_message_builder()
    messages = []

    messages.append(msg_builder.user().add_text("What is the weather in Paris?"))
    with tracking.set_tracker() as global_tracker:
        payload = APIPayload(messages=messages, tools=tools)
        llm_output1 = model(payload)

        prev_input = global_tracker.stats["input_tokens"]
        prev_output = global_tracker.stats["output_tokens"]
        prev_cost = global_tracker.stats["cost"]

        for tool_call in llm_output1.tool_calls:
            tool_call.response_text("It's sunny! 25°C")
        messages += [
            msg_builder.add_responded_tool_calls(llm_output1.tool_calls),
            msg_builder.user().add_text("What is the weather in Delhi?"),
        ]
        payload = APIPayload(messages=messages, tools=tools)
        llm_output2 = model(payload)
        delta_input = global_tracker.stats["input_tokens"] - prev_input
        delta_output = global_tracker.stats["output_tokens"] - prev_output
        delta_cost = global_tracker.stats["cost"] - prev_cost

    assert prev_input > 0, "Expected previous input tokens to be greater than 0"
    assert prev_output > 0, "Expected previous output tokens to be greater than 0"
    assert prev_cost > 0, "Expected previous cost value to be greater than 0"
    assert llm_output2.raw_response is not None
    assert (
        llm_output2.action == """get_weather(location='Delhi', unit='celsius')"""
    ), f"""Expected get_weather('Delhi') but got {llm_output2.action}"""
    assert delta_input > 0, "Expected new input tokens to be greater than 0"
    assert delta_output > 0, "Expected new output tokens to be greater than 0"
    assert delta_cost > 0, "Expected new cost value to be greater than 0"
    assert global_tracker.stats["input_tokens"] == prev_input + delta_input
    assert global_tracker.stats["output_tokens"] == prev_output + delta_output
    assert global_tracker.stats["cost"] == pytest.approx(prev_cost + delta_cost)

## Test multiaction
@pytest.mark.pricy
@pytest.mark.skipif(not os.getenv("OPENAI_API_KEY"), reason="Skipping as OpenAI API key not set")
def test_multi_action_tool_calls():
    """
    Test that the model can produce multiple tool calls in parallel.
    Uncomment commented lines to see the full behaviour of models and tool choices.
    """
    # test_config (setting name, BaseModelArgs, model_name, tools)
    tool_test_configs = [
        (
            "gpt-4.1-responses API",
            OpenAIResponseModelArgs,
            "gpt-4.1-2025-04-14",
            responses_api_tools,
        ),
        ("gpt-4.1-chat Completions API", OpenAIChatModelArgs, "gpt-4.1-2025-04-14", chat_api_tools),
        # ("claude-3", ClaudeResponseModelArgs, "claude-3-haiku-20240307", anthropic_tools),   # fails


        delta_cost = global_tracker.stats["cost"] - prev_cost

    assert prev_input > 0, "Expected previous input tokens to be greater than 0"
    assert prev_output > 0, "Expected previous output tokens to be greater than 0"
    assert prev_cost > 0, "Expected previous cost value to be greater than 0"
    assert llm_output2.raw_response is not None
    assert (
        llm_output2.action == """get_weather(location='Delhi', unit='celsius')"""
    ), f"""Expected get_weather('Delhi') but got {llm_output2.action}"""
    assert delta_input > 0, "Expected new input tokens to be greater than 0"
    assert delta_output > 0, "Expected new output tokens to be greater than 0"
    assert delta_cost > 0, "Expected new cost value to be greater than 0"
    assert global_tracker.stats["input_tokens"] == prev_input + delta_input
    assert global_tracker.stats["output_tokens"] == prev_output + delta_output
    assert global_tracker.stats["cost"] == pytest.approx(prev_cost + delta_cost)

## Test multiaction
@pytest.mark.pricy
@pytest.mark.skipif(not os.getenv("OPENAI_API_KEY"), reason="Skipping as OpenAI API key not set")
def test_multi_action_tool_calls():
    """
    Test that the model can produce multiple tool calls in parallel.
    Uncomment commented lines to see the full behaviour of models and tool choices.
    """
    # test_config (setting name, BaseModelArgs, model_name, tools)
    tool_test_configs = [
        (
            "gpt-4.1-responses API",
            OpenAIResponseModelArgs,
            "gpt-4.1-2025-04-14",
            responses_api_tools,
        ),
        ("gpt-4.1-chat Completions API", OpenAIChatModelArgs, "gpt-4.1-2025-04-14", chat_api_tools),
        # ("claude-3", ClaudeResponseModelArgs, "claude-3-haiku-20240307", anthropic_tools),   # fails
        # ("claude-3.7", ClaudeResponseModelArgs, "claude-3-7-sonnet-20250219", anthropic_tools), # fails
        ("claude-4-sonnet", ClaudeResponseModelArgs, "claude-sonnet-4-20250514", anthropic_tools),
        # add more models as needed
    ]

    def add_user_messages(msg_builder):
        return [
            msg_builder.user().add_text("What is the weather in Paris and Delhi?"),
            msg_builder.user().add_text("You must call multiple tools to achieve the task."),
        ]

    res_df = []

    for tool_choice in [
        # 'none',
        # 'required', # fails for Responses API
        # 'any',  # fails for Responses API
        "auto",
        # 'get_weather'
    ]:
        for name, llm_class, checkpoint_name, tools in tool_test_configs:
            print(name, "tool choice:", tool_choice, "\n", "**" * 10)
            model_args = llm_class(model_name=checkpoint_name, max_new_tokens=200, temperature=None)
            llm, msg_builder = model_args.make_model(), model_args.get_message_builder()
            messages = add_user_messages(msg_builder)
            if tool_choice == "get_weather":  # force a specific tool call
                response: LLMOutput = llm(
                    APIPayload(messages=messages, tools=tools, force_call_tool=tool_choice)
                )
            else:
                response: LLMOutput = llm(
                    APIPayload(messages=messages, tools=tools, tool_choice=tool_choice)
                )
                num_tool_calls = len(response.tool_calls) if response.tool_calls else 0
            res_df.append(
                {
                    "model": name,
                    "checkpoint": checkpoint_name,
                    "tool_choice": tool_choice,
                    "num_tool_calls": num_tool_calls,
                    "action": response.action,
                }
            )
            assert (
                num_tool_calls == 2
            ), f"Expected 2 tool calls, but got {num_tool_calls} for {name} with tool choice {tool_choice}"
        # import pandas as pd
        # print(pd.DataFrame(res_df))

EDGE_CASES = [
    # 1. Empty kwargs dict
    ("valid_function", {}, "valid_function()"),
    # 2. Kwargs with problematic string values (quotes, escapes, unicode)
    (
        "send_message",
        {
            "text": 'He said "Hello!" and used a backslash: \\',
            "unicode": "Café naïve résumé 🚀",
            "newlines": "Line1\nLine2\tTabbed",
        },
        "send_message(text='He said \"Hello!\" and used a backslash: \\\\', unicode='Café naïve résumé 🚀', newlines='Line1\\nLine2\\tTabbed')",
    ),
    # 3. Mixed types including problematic float values
    (
        "complex_call",


        {
            "infinity": float("inf"),
            "nan": float("nan"),
            "negative_zero": -0.0,
            "scientific": 1.23e-45,
        },
        "complex_call(infinity=inf, nan=nan, negative_zero=-0.0, scientific=1.23e-45)",
    ),
    # 4. Deeply nested structures that could stress repr()
    (
        "process_data",
        {
            "nested": {"level1": {"level2": {"level3": [1, 2, {"deep": True}]}}},
            "circular_ref_like": {"a": {"b": {"c": "back_to_start"}}},
        },
        "process_data(nested={'level1': {'level2': {'level3': [1, 2, {'deep': True}]}}}, circular_ref_like={'a': {'b': {'c': 'back_to_start'}}})",
    ),
]

def test_tool_call_to_python_code():
    from agentlab.llm.response_api import tool_call_to_python_code

    for edge_case in EDGE_CASES:
        func_name, kwargs, expected = edge_case
        result = tool_call_to_python_code(func_name, kwargs)
        print(result)
        assert result == expected, f"Expected {expected} but got {result}"

if __name__ == "__main__":
    test_tool_call_to_python_code()
    # test_openai_chat_completion_model_parse_and_cost()


def test_multi_action_tool_calls():
    """
    Test that the model can produce multiple tool calls in parallel.
    Uncomment commented lines to see the full behaviour of models and tool choices.
    """
    # test_config (setting name, BaseModelArgs, model_name, tools)
    tool_test_configs = [
        (
            "gpt-4.1-responses API",
            OpenAIResponseModelArgs,
            "gpt-4.1-2025-04-14",
            responses_api_tools,
        ),
        ("gpt-4.1-chat Completions API", OpenAIChatModelArgs, "gpt-4.1-2025-04-14", chat_api_tools),
        # ("claude-3", ClaudeResponseModelArgs, "claude-3-haiku-20240307", anthropic_tools),   # fails
        # ("claude-3.7", ClaudeResponseModelArgs, "claude-3-7-sonnet-20250219", anthropic_tools), # fails
        ("claude-4-sonnet", ClaudeResponseModelArgs, "claude-sonnet-4-20250514", anthropic_tools),
        # add more models as needed
    ]

    def add_user_messages(msg_builder):
        return [
            msg_builder.user().add_text("What is the weather in Paris and Delhi?"),
            msg_builder.user().add_text("You must call multiple tools to achieve the task."),
        ]

    res_df = []

    for tool_choice in [
        # 'none',
        # 'required', # fails for Responses API
        # 'any',  # fails for Responses API
        "auto",
        # 'get_weather'
    ]:
        for name, llm_class, checkpoint_name, tools in tool_test_configs:
            print(name, "tool choice:", tool_choice, "\n", "**" * 10)
            model_args = llm_class(model_name=checkpoint_name, max_new_tokens=200, temperature=None)
            llm, msg_builder = model_args.make_model(), model_args.get_message_builder()
            messages = add_user_messages(msg_builder)
            if tool_choice == "get_weather":  # force a specific tool call
                response: LLMOutput = llm(
                    APIPayload(messages=messages, tools=tools, force_call_tool=tool_choice)
                )
            else:


import os

import pytest

from agentlab.llm.chat_api import (
    AnthropicModelArgs,
    AzureModelArgs,
    OpenAIModelArgs,
    make_system_message,
    make_user_message,
)

# TODO(optimass): figure out a good model for all tests

if "AGENTLAB_LOCAL_TEST" in os.environ:
    skip_tests = os.environ["AGENTLAB_LOCAL_TEST"] != "1"
else:
    skip_tests = False

@pytest.mark.pricy
@pytest.mark.skipif(skip_tests, reason="Skipping on remote as Azure is pricy")
@pytest.mark.skipif(
    not os.getenv("AZURE_OPENAI_API_KEY"), reason="Skipping as Azure API key not set"
)
def test_api_model_args_azure():
    model_args = AzureModelArgs(
        model_name="gpt-4.1-nano",
        deployment_name="gpt-4.1-nano",
        max_total_tokens=8192,
        max_input_tokens=8192 - 512,
        max_new_tokens=512,
        temperature=1e-1,
    )
    model = model_args.make_model()

    messages = [
        make_system_message("You are an helpful virtual assistant"),
        make_user_message("Give the third prime number"),
    ]
    answer = model(messages)

    assert "5" in answer.get("content")

@pytest.mark.pricy
@pytest.mark.skipif(skip_tests, reason="Skipping on remote as Azure is pricy")
@pytest.mark.skipif(not os.getenv("OPENAI_API_KEY"), reason="Skipping as OpenAI API key not set")
def test_api_model_args_openai():
    model_args = OpenAIModelArgs(
        model_name="gpt-4o-mini",
        max_total_tokens=8192,
        max_input_tokens=8192 - 512,
        max_new_tokens=512,
        temperature=1e-1,
    )
    model = model_args.make_model()

    messages = [
        make_system_message("You are an helpful virtual assistant"),
        make_user_message("Give the third prime number"),
    ]
    answer = model(messages)

    assert "5" in answer.get("content")

@pytest.mark.pricy
@pytest.mark.skipif(skip_tests, reason="Skipping on remote as Anthropic is pricy")
@pytest.mark.skipif(
    not os.getenv("ANTHROPIC_API_KEY"), reason="Skipping as Anthropic API key not set"
)
def test_api_model_args_anthropic():
    model_args = AnthropicModelArgs(
        model_name="claude-3-haiku-20240307",
        max_total_tokens=8192,
        max_input_tokens=8192 - 512,
        max_new_tokens=512,
        temperature=1e-1,
    )
    model = model_args.make_model()

    messages = [
        make_system_message("You are an helpful virtual assistant"),
        make_user_message("Give the third prime number. Just the number, no explanation."),
    ]
    answer = model(messages)
    assert "5" in answer.get("content")

if __name__ == "__main__":
    test_api_model_args_anthropic()


    AzureModelArgs,
    OpenAIModelArgs,
    make_system_message,
    make_user_message,
)

# TODO(optimass): figure out a good model for all tests

if "AGENTLAB_LOCAL_TEST" in os.environ:
    skip_tests = os.environ["AGENTLAB_LOCAL_TEST"] != "1"
else:
    skip_tests = False

@pytest.mark.pricy
@pytest.mark.skipif(skip_tests, reason="Skipping on remote as Azure is pricy")
@pytest.mark.skipif(
    not os.getenv("AZURE_OPENAI_API_KEY"), reason="Skipping as Azure API key not set"
)
def test_api_model_args_azure():
    model_args = AzureModelArgs(
        model_name="gpt-4.1-nano",
        deployment_name="gpt-4.1-nano",
        max_total_tokens=8192,
        max_input_tokens=8192 - 512,
        max_new_tokens=512,
        temperature=1e-1,
    )
    model = model_args.make_model()

    messages = [
        make_system_message("You are an helpful virtual assistant"),
        make_user_message("Give the third prime number"),
    ]
    answer = model(messages)

    assert "5" in answer.get("content")

@pytest.mark.pricy
@pytest.mark.skipif(skip_tests, reason="Skipping on remote as Azure is pricy")
@pytest.mark.skipif(not os.getenv("OPENAI_API_KEY"), reason="Skipping as OpenAI API key not set")
def test_api_model_args_openai():
    model_args = OpenAIModelArgs(
        model_name="gpt-4o-mini",
        max_total_tokens=8192,
        max_input_tokens=8192 - 512,
        max_new_tokens=512,
        temperature=1e-1,
    )
    model = model_args.make_model()

    messages = [
        make_system_message("You are an helpful virtual assistant"),
        make_user_message("Give the third prime number"),
    ]
    answer = model(messages)


        deployment_name="gpt-4.1-nano",
        max_total_tokens=8192,
        max_input_tokens=8192 - 512,
        max_new_tokens=512,
        temperature=1e-1,
    )
    model = model_args.make_model()

    messages = [
        make_system_message("You are an helpful virtual assistant"),
        make_user_message("Give the third prime number"),
    ]
    answer = model(messages)

    assert "5" in answer.get("content")

@pytest.mark.pricy
@pytest.mark.skipif(skip_tests, reason="Skipping on remote as Azure is pricy")
@pytest.mark.skipif(not os.getenv("OPENAI_API_KEY"), reason="Skipping as OpenAI API key not set")
def test_api_model_args_openai():
    model_args = OpenAIModelArgs(
        model_name="gpt-4o-mini",
        max_total_tokens=8192,
        max_input_tokens=8192 - 512,
        max_new_tokens=512,
        temperature=1e-1,
    )
    model = model_args.make_model()

    messages = [
        make_system_message("You are an helpful virtual assistant"),
        make_user_message("Give the third prime number"),
    ]
    answer = model(messages)

    assert "5" in answer.get("content")

@pytest.mark.pricy
@pytest.mark.skipif(skip_tests, reason="Skipping on remote as Anthropic is pricy")
@pytest.mark.skipif(
    not os.getenv("ANTHROPIC_API_KEY"), reason="Skipping as Anthropic API key not set"
)
def test_api_model_args_anthropic():
    model_args = AnthropicModelArgs(
        model_name="claude-3-haiku-20240307",
        max_total_tokens=8192,
        max_input_tokens=8192 - 512,
        max_new_tokens=512,
        temperature=1e-1,
    )
    model = model_args.make_model()

    messages = [
        make_system_message("You are an helpful virtual assistant"),
        make_user_message("Give the third prime number. Just the number, no explanation."),


        max_input_tokens=8192 - 512,
        max_new_tokens=512,
        temperature=1e-1,
    )
    model = model_args.make_model()

    messages = [
        make_system_message("You are an helpful virtual assistant"),
        make_user_message("Give the third prime number"),
    ]
    answer = model(messages)

    assert "5" in answer.get("content")

@pytest.mark.pricy
@pytest.mark.skipif(skip_tests, reason="Skipping on remote as Anthropic is pricy")
@pytest.mark.skipif(
    not os.getenv("ANTHROPIC_API_KEY"), reason="Skipping as Anthropic API key not set"
)
def test_api_model_args_anthropic():
    model_args = AnthropicModelArgs(
        model_name="claude-3-haiku-20240307",
        max_total_tokens=8192,
        max_input_tokens=8192 - 512,
        max_new_tokens=512,
        temperature=1e-1,
    )
    model = model_args.make_model()

    messages = [
        make_system_message("You are an helpful virtual assistant"),
        make_user_message("Give the third prime number. Just the number, no explanation."),
    ]
    answer = model(messages)
    assert "5" in answer.get("content")

if __name__ == "__main__":
    test_api_model_args_anthropic()


import warnings
from typing import Literal
from unittest.mock import Mock

import httpx
import pytest
from openai import RateLimitError

from agentlab.llm import llm_utils
from agentlab.llm.chat_api import make_system_message

yaml_str = """Analysis:
This is the analysis

Summary: This is the summary

Confidence Score: 7
"""

def test_yaml_parser():
    ans, _, _ = llm_utils.yaml_parser(yaml_str)
    print(ans)
    assert ans["Analysis"] == "This is the analysis"
    assert ans["Summary"] == "This is the summary"
    assert ans["Confidence Score"] == 7

def test_truncate_tokens():
    text = "This is a simple test."
    truncated = llm_utils.truncate_tokens(text, max_tokens=3)
    assert truncated == "This is a"

def test_count_tokens():
    text = "This is a simple test."
    assert llm_utils.count_tokens(text) == 6

def test_json_parser():
    # Testing valid JSON
    message = '{"test": "Hello, World!"}'

    # deactivate warnings
    warnings.filterwarnings("ignore")

    value, valid, retry_message = llm_utils.json_parser(message)
    assert value == {"test": "Hello, World!"}
    assert valid == True
    assert retry_message == ""

    # Testing invalid JSON
    message = '{"test": "Hello, World!"'  # missing closing brace
    value, valid, retry_message = llm_utils.json_parser(message)
    assert value == {}
    assert valid == False
    assert len(retry_message) > 3

    # reactivate warnings
    warnings.filterwarnings("default")

def test_compress_string():
    text = """
This is a test
for paragraph.

This is a second test.
hola
This is a second test.

This is a test
for paragraph.
"""

    expected_output = """\
<definitions>
§-0:
This is a test
for paragraph.
¶-0:
This is a second test.
</definitions>
§-0
¶-0
hola
¶-0
§-0"""

    compressed_text = llm_utils.compress_string(text)
    assert compressed_text == expected_output

# Mock ChatOpenAI class
class MockChatOpenAI:
    def call(self, messages):
        return "mocked response"

    def __call__(self, messages):
        return self.call(messages)

def mock_parser(answer):
    if answer == "correct content":
        return "Parsed value"
    else:
        raise llm_utils.ParseError("Retry message")

def mock_rate_limit_error(message: str, status_code: Literal[429] = 429) -> RateLimitError:
    """
    Create a mocked instantiation of RateLimitError with a specified message and status code.

    Args:
        message (str): The error message.
        status_code (Literal[429]): The HTTP status code, default is 429 for rate limiting.

    Returns:
        RateLimitError: A mocked RateLimitError instance.
    """
    mock_response = Mock(spec=httpx.Response)
    mock_response.status_code = status_code
    mock_response.json.return_value = {"error": {"message": message}}
    mock_response.headers = {"x-request-id": "test-request-id"}  # Add headers attribute

    return RateLimitError(message, response=mock_response, body=mock_response.json())

# Test to ensure function stops retrying after reaching the max wait time
# def test_rate_limit_max_wait_time():
#     mock_chat = MockChatOpenAI()
#     mock_chat.call = Mock(
#         side_effect=mock_rate_limit_error("Rate limit reached. Please try again in 2s.")
#     )

#     with pytest.raises(RateLimitError):
#         llm_utils.retry(
#             mock_chat,
#             [],
#             n_retry=4,
#             parser=mock_parser,
#             rate_limit_max_wait_time=6,
#             min_retry_wait_time=1,
#         )

#     # The function should stop retrying after 2 attempts (6s each time, 12s total which is greater than the 10s max wait time)
#     assert mock_chat.call.call_count == 3

# def test_rate_limit_success():
#     mock_chat = MockChatOpenAI()
#     mock_chat.call = Mock(
#         side_effect=[
#             mock_rate_limit_error("Rate limit reached. Please try again in 2s."),
#             make_system_message("correct content"),
#         ]
#     )

#     result = llm_utils.retry(
#         mock_chat,
#         [],
#         n_retry=4,
#         parser=mock_parser,
#         rate_limit_max_wait_time=6,
#         min_retry_wait_time=1,
#     )

#     assert result == "Parsed value"
#     assert mock_chat.call.call_count == 2

# Mock a successful parser response to test function exit before max retries
def test_successful_parse_before_max_retries():
    mock_chat = MockChatOpenAI()

    # mock a chat that returns the wrong content the first 2 time, but the right
    # content  on the 3rd time
    mock_chat.call = Mock(
        side_effect=[
            make_system_message("wrong content"),
            make_system_message("wrong content"),
            make_system_message("correct content"),
        ]
    )

    result = llm_utils.retry(mock_chat, llm_utils.Discussion(), 5, mock_parser)

    assert result == "Parsed value"
    assert mock_chat.call.call_count == 3

def test_unsuccessful_parse_before_max_retries():
    mock_chat = MockChatOpenAI()

    # mock a chat that returns the wrong content the first 2 time, but the right
    # content  on the 3rd time
    mock_chat.call = Mock(
        side_effect=[
            make_system_message("wrong content"),
            make_system_message("wrong content"),
            make_system_message("correct content"),
        ]
    )
    with pytest.raises(llm_utils.ParseError):
        result = llm_utils.retry(mock_chat, llm_utils.Discussion(), 2, mock_parser)

    assert mock_chat.call.call_count == 2

def test_retry_parse_raises():
    mock_chat = MockChatOpenAI()
    mock_chat.call = Mock(return_value=make_system_message("mocked response"))
    parser_raises = Mock(side_effect=ValueError("Parser error"))

    with pytest.raises(ValueError):
        llm_utils.retry(mock_chat, llm_utils.Discussion(), 3, parser_raises)

def test_extract_code_blocks():
    text = """\
This is some text.
```python
def hello_world():
    print("Hello, world!")
```
Some more text.
```
More code without a language.
```
Another block of code:
```javascript
console.log("Hello, world!");
```
An inline code block ```click()```
"""

    expected_output = [
        ("python", 'def hello_world():\n    print("Hello, world!")'),
        ("", "More code without a language."),
        ("javascript", 'console.log("Hello, world!");'),
        ("", "click()"),
    ]

    assert llm_utils.extract_code_blocks(text) == expected_output

def test_message_merge_only_text():
    content = [
        {"type": "text", "text": "Hello, world!"},
        {"type": "text", "text": "This is a test."},
    ]
    message = llm_utils.BaseMessage(role="system", content=content)
    message.merge()
    assert message["content"] == "Hello, world!\nThis is a test."

def test_message_merge_text_image():
    content = [
        {"type": "text", "text": "Hello, world!"},
        {"type": "text", "text": "This is a test."},
        {"type": "image_url", "image_url": "this is a base64 image"},
        {"type": "text", "text": "This is another test."},
        {"type": "text", "text": "Goodbye, world!"},
    ]
    message = llm_utils.BaseMessage(role="system", content=content)
    message.merge()
    assert len(message["content"]) == 3
    assert message["content"][0]["text"] == "Hello, world!\nThis is a test."
    assert message["content"][1]["image_url"] == "this is a base64 image"
    assert message["content"][2]["text"] == "This is another test.\nGoodbye, world!"

if __name__ == "__main__":
    # test_retry_parallel()
    # test_rate_limit_max_wait_time()
    # test_successful_parse_before_max_retries()
    # test_unsuccessful_parse_before_max_retries()
    # test_extract_code_blocks()
    # test_message_merge_only_text()
    test_message_merge_text_image()


import warnings
from typing import Literal
from unittest.mock import Mock

import httpx
import pytest
from openai import RateLimitError

from agentlab.llm import llm_utils
from agentlab.llm.chat_api import make_system_message

yaml_str = """Analysis:
This is the analysis

Summary: This is the summary

Confidence Score: 7
"""

def test_yaml_parser():
    ans, _, _ = llm_utils.yaml_parser(yaml_str)
    print(ans)
    assert ans["Analysis"] == "This is the analysis"
    assert ans["Summary"] == "This is the summary"
    assert ans["Confidence Score"] == 7

def test_truncate_tokens():
    text = "This is a simple test."
    truncated = llm_utils.truncate_tokens(text, max_tokens=3)
    assert truncated == "This is a"

def test_count_tokens():
    text = "This is a simple test."
    assert llm_utils.count_tokens(text) == 6

def test_json_parser():
    # Testing valid JSON
    message = '{"test": "Hello, World!"}'

    # deactivate warnings
    warnings.filterwarnings("ignore")


from agentlab.llm import llm_utils
from agentlab.llm.chat_api import make_system_message

yaml_str = """Analysis:
This is the analysis

Summary: This is the summary

Confidence Score: 7
"""

def test_yaml_parser():
    ans, _, _ = llm_utils.yaml_parser(yaml_str)
    print(ans)
    assert ans["Analysis"] == "This is the analysis"
    assert ans["Summary"] == "This is the summary"
    assert ans["Confidence Score"] == 7

def test_truncate_tokens():
    text = "This is a simple test."
    truncated = llm_utils.truncate_tokens(text, max_tokens=3)
    assert truncated == "This is a"

def test_count_tokens():
    text = "This is a simple test."
    assert llm_utils.count_tokens(text) == 6

def test_json_parser():
    # Testing valid JSON
    message = '{"test": "Hello, World!"}'

    # deactivate warnings
    warnings.filterwarnings("ignore")

    value, valid, retry_message = llm_utils.json_parser(message)
    assert value == {"test": "Hello, World!"}
    assert valid == True
    assert retry_message == ""

    # Testing invalid JSON


Summary: This is the summary

Confidence Score: 7
"""

def test_yaml_parser():
    ans, _, _ = llm_utils.yaml_parser(yaml_str)
    print(ans)
    assert ans["Analysis"] == "This is the analysis"
    assert ans["Summary"] == "This is the summary"
    assert ans["Confidence Score"] == 7

def test_truncate_tokens():
    text = "This is a simple test."
    truncated = llm_utils.truncate_tokens(text, max_tokens=3)
    assert truncated == "This is a"

def test_count_tokens():
    text = "This is a simple test."
    assert llm_utils.count_tokens(text) == 6

def test_json_parser():
    # Testing valid JSON
    message = '{"test": "Hello, World!"}'

    # deactivate warnings
    warnings.filterwarnings("ignore")

    value, valid, retry_message = llm_utils.json_parser(message)
    assert value == {"test": "Hello, World!"}
    assert valid == True
    assert retry_message == ""

    # Testing invalid JSON
    message = '{"test": "Hello, World!"'  # missing closing brace
    value, valid, retry_message = llm_utils.json_parser(message)
    assert value == {}
    assert valid == False
    assert len(retry_message) > 3


def test_yaml_parser():
    ans, _, _ = llm_utils.yaml_parser(yaml_str)
    print(ans)
    assert ans["Analysis"] == "This is the analysis"
    assert ans["Summary"] == "This is the summary"
    assert ans["Confidence Score"] == 7

def test_truncate_tokens():
    text = "This is a simple test."
    truncated = llm_utils.truncate_tokens(text, max_tokens=3)
    assert truncated == "This is a"

def test_count_tokens():
    text = "This is a simple test."
    assert llm_utils.count_tokens(text) == 6

def test_json_parser():
    # Testing valid JSON
    message = '{"test": "Hello, World!"}'

    # deactivate warnings
    warnings.filterwarnings("ignore")

    value, valid, retry_message = llm_utils.json_parser(message)
    assert value == {"test": "Hello, World!"}
    assert valid == True
    assert retry_message == ""

    # Testing invalid JSON
    message = '{"test": "Hello, World!"'  # missing closing brace
    value, valid, retry_message = llm_utils.json_parser(message)
    assert value == {}
    assert valid == False
    assert len(retry_message) > 3

    # reactivate warnings
    warnings.filterwarnings("default")

def test_compress_string():
    text = """
This is a test
for paragraph.

This is a second test.
hola
This is a second test.

This is a test
for paragraph.
"""

    expected_output = """\
<definitions>
§-0:
This is a test
for paragraph.


    # deactivate warnings
    warnings.filterwarnings("ignore")

    value, valid, retry_message = llm_utils.json_parser(message)
    assert value == {"test": "Hello, World!"}
    assert valid == True
    assert retry_message == ""

    # Testing invalid JSON
    message = '{"test": "Hello, World!"'  # missing closing brace
    value, valid, retry_message = llm_utils.json_parser(message)
    assert value == {}
    assert valid == False
    assert len(retry_message) > 3

    # reactivate warnings
    warnings.filterwarnings("default")

def test_compress_string():
    text = """
This is a test
for paragraph.

This is a second test.
hola
This is a second test.

This is a test
for paragraph.
"""

    expected_output = """\
<definitions>
§-0:
This is a test
for paragraph.
¶-0:
This is a second test.
</definitions>
§-0
¶-0
hola
¶-0
§-0"""

    compressed_text = llm_utils.compress_string(text)
    assert compressed_text == expected_output

# Mock ChatOpenAI class
class MockChatOpenAI:
    def call(self, messages):
        return "mocked response"

    def __call__(self, messages):
        return self.call(messages)

def mock_parser(answer):
    if answer == "correct content":
        return "Parsed value"
    else:
        raise llm_utils.ParseError("Retry message")

def mock_rate_limit_error(message: str, status_code: Literal[429] = 429) -> RateLimitError:
    """


MockChatOpenAI

    expected_output = """\
<definitions>
§-0:
This is a test
for paragraph.
¶-0:
This is a second test.
</definitions>
§-0
¶-0
hola
¶-0
§-0"""

    compressed_text = llm_utils.compress_string(text)
    assert compressed_text == expected_output

# Mock ChatOpenAI class
class MockChatOpenAI:
    def call(self, messages):
        return "mocked response"

    def __call__(self, messages):
        return self.call(messages)

def mock_parser(answer):
    if answer == "correct content":
        return "Parsed value"
    else:
        raise llm_utils.ParseError("Retry message")

def mock_rate_limit_error(message: str, status_code: Literal[429] = 429) -> RateLimitError:
    """
    Create a mocked instantiation of RateLimitError with a specified message and status code.

    Args:
        message (str): The error message.
        status_code (Literal[429]): The HTTP status code, default is 429 for rate limiting.

    Returns:
        RateLimitError: A mocked RateLimitError instance.
    """


</definitions>
§-0
¶-0
hola
¶-0
§-0"""

    compressed_text = llm_utils.compress_string(text)
    assert compressed_text == expected_output

# Mock ChatOpenAI class
class MockChatOpenAI:
    def call(self, messages):
        return "mocked response"

    def __call__(self, messages):
        return self.call(messages)

def mock_parser(answer):
    if answer == "correct content":
        return "Parsed value"
    else:
        raise llm_utils.ParseError("Retry message")

def mock_rate_limit_error(message: str, status_code: Literal[429] = 429) -> RateLimitError:
    """
    Create a mocked instantiation of RateLimitError with a specified message and status code.

    Args:
        message (str): The error message.
        status_code (Literal[429]): The HTTP status code, default is 429 for rate limiting.

    Returns:
        RateLimitError: A mocked RateLimitError instance.
    """
    mock_response = Mock(spec=httpx.Response)
    mock_response.status_code = status_code
    mock_response.json.return_value = {"error": {"message": message}}
    mock_response.headers = {"x-request-id": "test-request-id"}  # Add headers attribute

    return RateLimitError(message, response=mock_response, body=mock_response.json())


    compressed_text = llm_utils.compress_string(text)
    assert compressed_text == expected_output

# Mock ChatOpenAI class
class MockChatOpenAI:
    def call(self, messages):
        return "mocked response"

    def __call__(self, messages):
        return self.call(messages)

def mock_parser(answer):
    if answer == "correct content":
        return "Parsed value"
    else:
        raise llm_utils.ParseError("Retry message")

def mock_rate_limit_error(message: str, status_code: Literal[429] = 429) -> RateLimitError:
    """
    Create a mocked instantiation of RateLimitError with a specified message and status code.

    Args:
        message (str): The error message.
        status_code (Literal[429]): The HTTP status code, default is 429 for rate limiting.

    Returns:
        RateLimitError: A mocked RateLimitError instance.
    """
    mock_response = Mock(spec=httpx.Response)
    mock_response.status_code = status_code
    mock_response.json.return_value = {"error": {"message": message}}
    mock_response.headers = {"x-request-id": "test-request-id"}  # Add headers attribute

    return RateLimitError(message, response=mock_response, body=mock_response.json())

# Test to ensure function stops retrying after reaching the max wait time
# def test_rate_limit_max_wait_time():
#     mock_chat = MockChatOpenAI()
#     mock_chat.call = Mock(
#         side_effect=mock_rate_limit_error("Rate limit reached. Please try again in 2s.")
#     )

#     with pytest.raises(RateLimitError):
#         llm_utils.retry(
#             mock_chat,
#             [],
#             n_retry=4,
#             parser=mock_parser,
#             rate_limit_max_wait_time=6,
#             min_retry_wait_time=1,
#         )

#     # The function should stop retrying after 2 attempts (6s each time, 12s total which is greater than the 10s max wait time)


#         side_effect=[
#             mock_rate_limit_error("Rate limit reached. Please try again in 2s."),
#             make_system_message("correct content"),
#         ]
#     )

#     result = llm_utils.retry(
#         mock_chat,
#         [],
#         n_retry=4,
#         parser=mock_parser,
#         rate_limit_max_wait_time=6,
#         min_retry_wait_time=1,
#     )

#     assert result == "Parsed value"
#     assert mock_chat.call.call_count == 2

# Mock a successful parser response to test function exit before max retries
def test_successful_parse_before_max_retries():
    mock_chat = MockChatOpenAI()

    # mock a chat that returns the wrong content the first 2 time, but the right
    # content  on the 3rd time
    mock_chat.call = Mock(
        side_effect=[
            make_system_message("wrong content"),
            make_system_message("wrong content"),
            make_system_message("correct content"),
        ]
    )

    result = llm_utils.retry(mock_chat, llm_utils.Discussion(), 5, mock_parser)

    assert result == "Parsed value"
    assert mock_chat.call.call_count == 3

def test_unsuccessful_parse_before_max_retries():
    mock_chat = MockChatOpenAI()

    # mock a chat that returns the wrong content the first 2 time, but the right
    # content  on the 3rd time
    mock_chat.call = Mock(
        side_effect=[
            make_system_message("wrong content"),
            make_system_message("wrong content"),
            make_system_message("correct content"),
        ]
    )
    with pytest.raises(llm_utils.ParseError):
        result = llm_utils.retry(mock_chat, llm_utils.Discussion(), 2, mock_parser)

    assert mock_chat.call.call_count == 2


# Mock a successful parser response to test function exit before max retries
def test_successful_parse_before_max_retries():
    mock_chat = MockChatOpenAI()

    # mock a chat that returns the wrong content the first 2 time, but the right
    # content  on the 3rd time
    mock_chat.call = Mock(
        side_effect=[
            make_system_message("wrong content"),
            make_system_message("wrong content"),
            make_system_message("correct content"),
        ]
    )

    result = llm_utils.retry(mock_chat, llm_utils.Discussion(), 5, mock_parser)

    assert result == "Parsed value"
    assert mock_chat.call.call_count == 3

def test_unsuccessful_parse_before_max_retries():
    mock_chat = MockChatOpenAI()

    # mock a chat that returns the wrong content the first 2 time, but the right
    # content  on the 3rd time
    mock_chat.call = Mock(
        side_effect=[
            make_system_message("wrong content"),
            make_system_message("wrong content"),
            make_system_message("correct content"),
        ]
    )
    with pytest.raises(llm_utils.ParseError):
        result = llm_utils.retry(mock_chat, llm_utils.Discussion(), 2, mock_parser)

    assert mock_chat.call.call_count == 2

def test_retry_parse_raises():
    mock_chat = MockChatOpenAI()
    mock_chat.call = Mock(return_value=make_system_message("mocked response"))
    parser_raises = Mock(side_effect=ValueError("Parser error"))

    with pytest.raises(ValueError):
        llm_utils.retry(mock_chat, llm_utils.Discussion(), 3, parser_raises)

def test_extract_code_blocks():
    text = """\
This is some text.
```python
def hello_world():
    print("Hello, world!")
```
Some more text.
```


def test_unsuccessful_parse_before_max_retries():
    mock_chat = MockChatOpenAI()

    # mock a chat that returns the wrong content the first 2 time, but the right
    # content  on the 3rd time
    mock_chat.call = Mock(
        side_effect=[
            make_system_message("wrong content"),
            make_system_message("wrong content"),
            make_system_message("correct content"),
        ]
    )
    with pytest.raises(llm_utils.ParseError):
        result = llm_utils.retry(mock_chat, llm_utils.Discussion(), 2, mock_parser)

    assert mock_chat.call.call_count == 2

def test_retry_parse_raises():
    mock_chat = MockChatOpenAI()
    mock_chat.call = Mock(return_value=make_system_message("mocked response"))
    parser_raises = Mock(side_effect=ValueError("Parser error"))

    with pytest.raises(ValueError):
        llm_utils.retry(mock_chat, llm_utils.Discussion(), 3, parser_raises)

def test_extract_code_blocks():
    text = """\
This is some text.
```python
def hello_world():
    print("Hello, world!")
```
Some more text.
```
More code without a language.
```
Another block of code:
```javascript
console.log("Hello, world!");
```
An inline code block ```click()```
"""


            make_system_message("wrong content"),
            make_system_message("wrong content"),
            make_system_message("correct content"),
        ]
    )
    with pytest.raises(llm_utils.ParseError):
        result = llm_utils.retry(mock_chat, llm_utils.Discussion(), 2, mock_parser)

    assert mock_chat.call.call_count == 2

def test_retry_parse_raises():
    mock_chat = MockChatOpenAI()
    mock_chat.call = Mock(return_value=make_system_message("mocked response"))
    parser_raises = Mock(side_effect=ValueError("Parser error"))

    with pytest.raises(ValueError):
        llm_utils.retry(mock_chat, llm_utils.Discussion(), 3, parser_raises)

def test_extract_code_blocks():
    text = """\
This is some text.
```python
def hello_world():
    print("Hello, world!")
```
Some more text.
```
More code without a language.
```
Another block of code:
```javascript
console.log("Hello, world!");
```
An inline code block ```click()```
"""

    expected_output = [
        ("python", 'def hello_world():\n    print("Hello, world!")'),
        ("", "More code without a language."),
        ("javascript", 'console.log("Hello, world!");'),
        ("", "click()"),
    ]

    assert llm_utils.extract_code_blocks(text) == expected_output

def test_message_merge_only_text():
    content = [
        {"type": "text", "text": "Hello, world!"},
        {"type": "text", "text": "This is a test."},
    ]
    message = llm_utils.BaseMessage(role="system", content=content)
    message.merge()
    assert message["content"] == "Hello, world!\nThis is a test."

def test_message_merge_text_image():
    content = [
        {"type": "text", "text": "Hello, world!"},
        {"type": "text", "text": "This is a test."},
        {"type": "image_url", "image_url": "this is a base64 image"},
        {"type": "text", "text": "This is another test."},
        {"type": "text", "text": "Goodbye, world!"},
    ]


```
More code without a language.
```
Another block of code:
```javascript
console.log("Hello, world!");
```
An inline code block ```click()```
"""

    expected_output = [
        ("python", 'def hello_world():\n    print("Hello, world!")'),
        ("", "More code without a language."),
        ("javascript", 'console.log("Hello, world!");'),
        ("", "click()"),
    ]

    assert llm_utils.extract_code_blocks(text) == expected_output

def test_message_merge_only_text():
    content = [
        {"type": "text", "text": "Hello, world!"},
        {"type": "text", "text": "This is a test."},
    ]
    message = llm_utils.BaseMessage(role="system", content=content)
    message.merge()
    assert message["content"] == "Hello, world!\nThis is a test."

def test_message_merge_text_image():
    content = [
        {"type": "text", "text": "Hello, world!"},
        {"type": "text", "text": "This is a test."},
        {"type": "image_url", "image_url": "this is a base64 image"},
        {"type": "text", "text": "This is another test."},
        {"type": "text", "text": "Goodbye, world!"},
    ]
    message = llm_utils.BaseMessage(role="system", content=content)
    message.merge()
    assert len(message["content"]) == 3
    assert message["content"][0]["text"] == "Hello, world!\nThis is a test."
    assert message["content"][1]["image_url"] == "this is a base64 image"
    assert message["content"][2]["text"] == "This is another test.\nGoodbye, world!"

if __name__ == "__main__":
    # test_retry_parallel()


    expected_output = [
        ("python", 'def hello_world():\n    print("Hello, world!")'),
        ("", "More code without a language."),
        ("javascript", 'console.log("Hello, world!");'),
        ("", "click()"),
    ]

    assert llm_utils.extract_code_blocks(text) == expected_output

def test_message_merge_only_text():
    content = [
        {"type": "text", "text": "Hello, world!"},
        {"type": "text", "text": "This is a test."},
    ]
    message = llm_utils.BaseMessage(role="system", content=content)
    message.merge()
    assert message["content"] == "Hello, world!\nThis is a test."

def test_message_merge_text_image():
    content = [
        {"type": "text", "text": "Hello, world!"},
        {"type": "text", "text": "This is a test."},
        {"type": "image_url", "image_url": "this is a base64 image"},
        {"type": "text", "text": "This is another test."},
        {"type": "text", "text": "Goodbye, world!"},
    ]
    message = llm_utils.BaseMessage(role="system", content=content)
    message.merge()
    assert len(message["content"]) == 3
    assert message["content"][0]["text"] == "Hello, world!\nThis is a test."
    assert message["content"][1]["image_url"] == "this is a base64 image"
    assert message["content"][2]["text"] == "This is another test.\nGoodbye, world!"

if __name__ == "__main__":
    # test_retry_parallel()
    # test_rate_limit_max_wait_time()
    # test_successful_parse_before_max_retries()
    # test_unsuccessful_parse_before_max_retries()
    # test_extract_code_blocks()
    # test_message_merge_only_text()
    test_message_merge_text_image()


call

    expected_output = """\
<definitions>
§-0:
This is a test
for paragraph.
¶-0:
This is a second test.
</definitions>
§-0
¶-0
hola
¶-0
§-0"""

    compressed_text = llm_utils.compress_string(text)
    assert compressed_text == expected_output

# Mock ChatOpenAI class
class MockChatOpenAI:
    def call(self, messages):
        return "mocked response"

    def __call__(self, messages):
        return self.call(messages)

def mock_parser(answer):
    if answer == "correct content":
        return "Parsed value"
    else:
        raise llm_utils.ParseError("Retry message")

def mock_rate_limit_error(message: str, status_code: Literal[429] = 429) -> RateLimitError:
    """
    Create a mocked instantiation of RateLimitError with a specified message and status code.

    Args:
        message (str): The error message.
        status_code (Literal[429]): The HTTP status code, default is 429 for rate limiting.


This is a test
for paragraph.
¶-0:
This is a second test.
</definitions>
§-0
¶-0
hola
¶-0
§-0"""

    compressed_text = llm_utils.compress_string(text)
    assert compressed_text == expected_output

# Mock ChatOpenAI class
class MockChatOpenAI:
    def call(self, messages):
        return "mocked response"

    def __call__(self, messages):
        return self.call(messages)

def mock_parser(answer):
    if answer == "correct content":
        return "Parsed value"
    else:
        raise llm_utils.ParseError("Retry message")

def mock_rate_limit_error(message: str, status_code: Literal[429] = 429) -> RateLimitError:
    """
    Create a mocked instantiation of RateLimitError with a specified message and status code.

    Args:
        message (str): The error message.
        status_code (Literal[429]): The HTTP status code, default is 429 for rate limiting.

    Returns:
        RateLimitError: A mocked RateLimitError instance.
    """


import math
import tempfile
from pathlib import Path

import pytest

from agentlab.agents.generic_agent.agent_configs import FLAGS_GPT_3_5, AGENT_4o_MINI
from agentlab.agents.generic_agent.generic_agent import GenericAgentArgs
from agentlab.analyze import inspect_results
from agentlab.experiments.launch_exp import (
    find_incomplete,
    non_dummy_count,
    run_experiments,
)
from agentlab.experiments.loop import EnvArgs, ExpArgs
from agentlab.experiments.study import Study
from agentlab.llm.chat_api import CheatMiniWoBLLMArgs

def test_relaunch_study():
    study_dir = Path(__file__).parent.parent / "data" / "test_study"
    exp_args_list = find_incomplete(study_dir, include_errors=False)

    assert non_dummy_count(exp_args_list) == 1
    assert exp_args_list[0].env_args.task_name == "miniwob.ascending-numbers"

    exp_args_list = find_incomplete(study_dir, include_errors=True)

    assert non_dummy_count(exp_args_list) == 2

def _test_launch_system(backend="ray", cause_timeout=False):
    if cause_timeout:
        wait_time = 10
        avg_step_timeout = 0.5
    else:
        wait_time = 0
        avg_step_timeout = 10

    exp_args_list = []
    for seed in range(3):
        exp_args_list.append(
            ExpArgs(
                agent_args=GenericAgentArgs(
                    chat_model_args=CheatMiniWoBLLMArgs(wait_time=wait_time),
                    flags=FLAGS_GPT_3_5,
                ),
                env_args=EnvArgs(task_name="miniwob.click-test", task_seed=seed, max_steps=5),
            )
        )

    with tempfile.TemporaryDirectory() as tmp_dir:
        study_dir = Path(tmp_dir) / "generic_agent_test"
        run_experiments(
            n_jobs=2,
            exp_args_list=exp_args_list,
            study_dir=study_dir,
            parallel_backend=backend,
            avg_step_timeout=avg_step_timeout,
        )

        results_df = inspect_results.load_result_df(study_dir, progress_fn=None)
        assert len(results_df) == len(exp_args_list)

        for _, row in results_df.iterrows():
            if row.stack_trace is not None:
                print(row.stack_trace)
            if cause_timeout:
                # assert row.err_msg is not None
                assert math.isnan(row.cum_reward) or row.cum_reward == 0
            else:
                assert row.err_msg is None
                assert row.cum_reward == 1.0

        study_summary = inspect_results.summarize_study(results_df)
        assert len(study_summary) == 1
        assert study_summary.std_err.iloc[0] == 0

        if not cause_timeout:
            assert study_summary.n_completed.iloc[0] == "3/3"
            assert study_summary.avg_reward.iloc[0] == 1.0

def test_launch_system_joblib():
    _test_launch_system(backend="joblib")

def test_launch_system_sequntial():
    _test_launch_system(backend="sequential")

def test_launch_system_ray():
    _test_launch_system(backend="ray")

@pytest.mark.pricy
def test_timeout_ray():
    _test_launch_system(backend="ray", cause_timeout=True)

@pytest.mark.pricy
def test_4o_mini_on_miniwob_tiny_test():
    """Run with `pytest -m pricy`."""
    with tempfile.TemporaryDirectory() as tmp_dir:
        study = Study(agent_args=[AGENT_4o_MINI], benchmark="miniwob_tiny_test", dir=tmp_dir)

        study.run(n_jobs=4)

        results_df = inspect_results.load_result_df(study.dir, progress_fn=None)

        for row in results_df.iterrows():
            if row[1].err_msg:
                print(row[1].err_msg)
                print(row[1].stack_trace)

        assert len(results_df) == len(study.exp_args_list)
        summary = inspect_results.summarize_study(results_df)
        print(summary)
        assert len(summary) == 1
        reward = summary.avg_reward.iloc[0]
        assert reward == 1.0

if __name__ == "__main__":
    test_timeout_ray()
    # test_4o_mini_on_miniwob_tiny_test()
    # test_launch_system_ray()
    # test_launch_system_sequntial()


import math
import tempfile
from pathlib import Path

import pytest

from agentlab.agents.generic_agent.agent_configs import FLAGS_GPT_3_5, AGENT_4o_MINI
from agentlab.agents.generic_agent.generic_agent import GenericAgentArgs
from agentlab.analyze import inspect_results
from agentlab.experiments.launch_exp import (
    find_incomplete,
    non_dummy_count,
    run_experiments,
)
from agentlab.experiments.loop import EnvArgs, ExpArgs
from agentlab.experiments.study import Study
from agentlab.llm.chat_api import CheatMiniWoBLLMArgs

def test_relaunch_study():
    study_dir = Path(__file__).parent.parent / "data" / "test_study"
    exp_args_list = find_incomplete(study_dir, include_errors=False)

    assert non_dummy_count(exp_args_list) == 1
    assert exp_args_list[0].env_args.task_name == "miniwob.ascending-numbers"

    exp_args_list = find_incomplete(study_dir, include_errors=True)

    assert non_dummy_count(exp_args_list) == 2

def _test_launch_system(backend="ray", cause_timeout=False):
    if cause_timeout:
        wait_time = 10
        avg_step_timeout = 0.5
    else:
        wait_time = 0
        avg_step_timeout = 10

    exp_args_list = []
    for seed in range(3):
        exp_args_list.append(
            ExpArgs(
                agent_args=GenericAgentArgs(
                    chat_model_args=CheatMiniWoBLLMArgs(wait_time=wait_time),
                    flags=FLAGS_GPT_3_5,
                ),
                env_args=EnvArgs(task_name="miniwob.click-test", task_seed=seed, max_steps=5),
            )


    non_dummy_count,
    run_experiments,
)
from agentlab.experiments.loop import EnvArgs, ExpArgs
from agentlab.experiments.study import Study
from agentlab.llm.chat_api import CheatMiniWoBLLMArgs

def test_relaunch_study():
    study_dir = Path(__file__).parent.parent / "data" / "test_study"
    exp_args_list = find_incomplete(study_dir, include_errors=False)

    assert non_dummy_count(exp_args_list) == 1
    assert exp_args_list[0].env_args.task_name == "miniwob.ascending-numbers"

    exp_args_list = find_incomplete(study_dir, include_errors=True)

    assert non_dummy_count(exp_args_list) == 2

def _test_launch_system(backend="ray", cause_timeout=False):
    if cause_timeout:
        wait_time = 10
        avg_step_timeout = 0.5
    else:
        wait_time = 0
        avg_step_timeout = 10

    exp_args_list = []
    for seed in range(3):
        exp_args_list.append(
            ExpArgs(
                agent_args=GenericAgentArgs(
                    chat_model_args=CheatMiniWoBLLMArgs(wait_time=wait_time),
                    flags=FLAGS_GPT_3_5,
                ),
                env_args=EnvArgs(task_name="miniwob.click-test", task_seed=seed, max_steps=5),
            )
        )

    with tempfile.TemporaryDirectory() as tmp_dir:
        study_dir = Path(tmp_dir) / "generic_agent_test"
        run_experiments(
            n_jobs=2,
            exp_args_list=exp_args_list,
            study_dir=study_dir,
            parallel_backend=backend,
            avg_step_timeout=avg_step_timeout,
        )

        results_df = inspect_results.load_result_df(study_dir, progress_fn=None)
        assert len(results_df) == len(exp_args_list)

        for _, row in results_df.iterrows():
            if row.stack_trace is not None:
                print(row.stack_trace)
            if cause_timeout:
                # assert row.err_msg is not None
                assert math.isnan(row.cum_reward) or row.cum_reward == 0
            else:
                assert row.err_msg is None
                assert row.cum_reward == 1.0

        study_summary = inspect_results.summarize_study(results_df)
        assert len(study_summary) == 1
        assert study_summary.std_err.iloc[0] == 0

        if not cause_timeout:
            assert study_summary.n_completed.iloc[0] == "3/3"
            assert study_summary.avg_reward.iloc[0] == 1.0

def test_launch_system_joblib():
    _test_launch_system(backend="joblib")

def test_launch_system_sequntial():
    _test_launch_system(backend="sequential")

def test_launch_system_ray():
    _test_launch_system(backend="ray")

@pytest.mark.pricy
def test_timeout_ray():
    _test_launch_system(backend="ray", cause_timeout=True)

@pytest.mark.pricy


        for _, row in results_df.iterrows():
            if row.stack_trace is not None:
                print(row.stack_trace)
            if cause_timeout:
                # assert row.err_msg is not None
                assert math.isnan(row.cum_reward) or row.cum_reward == 0
            else:
                assert row.err_msg is None
                assert row.cum_reward == 1.0

        study_summary = inspect_results.summarize_study(results_df)
        assert len(study_summary) == 1
        assert study_summary.std_err.iloc[0] == 0

        if not cause_timeout:
            assert study_summary.n_completed.iloc[0] == "3/3"
            assert study_summary.avg_reward.iloc[0] == 1.0

def test_launch_system_joblib():
    _test_launch_system(backend="joblib")

def test_launch_system_sequntial():
    _test_launch_system(backend="sequential")

def test_launch_system_ray():
    _test_launch_system(backend="ray")

@pytest.mark.pricy
def test_timeout_ray():
    _test_launch_system(backend="ray", cause_timeout=True)

@pytest.mark.pricy
def test_4o_mini_on_miniwob_tiny_test():
    """Run with `pytest -m pricy`."""
    with tempfile.TemporaryDirectory() as tmp_dir:
        study = Study(agent_args=[AGENT_4o_MINI], benchmark="miniwob_tiny_test", dir=tmp_dir)


            if cause_timeout:
                # assert row.err_msg is not None
                assert math.isnan(row.cum_reward) or row.cum_reward == 0
            else:
                assert row.err_msg is None
                assert row.cum_reward == 1.0

        study_summary = inspect_results.summarize_study(results_df)
        assert len(study_summary) == 1
        assert study_summary.std_err.iloc[0] == 0

        if not cause_timeout:
            assert study_summary.n_completed.iloc[0] == "3/3"
            assert study_summary.avg_reward.iloc[0] == 1.0

def test_launch_system_joblib():
    _test_launch_system(backend="joblib")

def test_launch_system_sequntial():
    _test_launch_system(backend="sequential")

def test_launch_system_ray():
    _test_launch_system(backend="ray")

@pytest.mark.pricy
def test_timeout_ray():
    _test_launch_system(backend="ray", cause_timeout=True)

@pytest.mark.pricy
def test_4o_mini_on_miniwob_tiny_test():
    """Run with `pytest -m pricy`."""
    with tempfile.TemporaryDirectory() as tmp_dir:
        study = Study(agent_args=[AGENT_4o_MINI], benchmark="miniwob_tiny_test", dir=tmp_dir)

        study.run(n_jobs=4)

        results_df = inspect_results.load_result_df(study.dir, progress_fn=None)


                assert row.err_msg is None
                assert row.cum_reward == 1.0

        study_summary = inspect_results.summarize_study(results_df)
        assert len(study_summary) == 1
        assert study_summary.std_err.iloc[0] == 0

        if not cause_timeout:
            assert study_summary.n_completed.iloc[0] == "3/3"
            assert study_summary.avg_reward.iloc[0] == 1.0

def test_launch_system_joblib():
    _test_launch_system(backend="joblib")

def test_launch_system_sequntial():
    _test_launch_system(backend="sequential")

def test_launch_system_ray():
    _test_launch_system(backend="ray")

@pytest.mark.pricy
def test_timeout_ray():
    _test_launch_system(backend="ray", cause_timeout=True)

@pytest.mark.pricy
def test_4o_mini_on_miniwob_tiny_test():
    """Run with `pytest -m pricy`."""
    with tempfile.TemporaryDirectory() as tmp_dir:
        study = Study(agent_args=[AGENT_4o_MINI], benchmark="miniwob_tiny_test", dir=tmp_dir)

        study.run(n_jobs=4)

        results_df = inspect_results.load_result_df(study.dir, progress_fn=None)

        for row in results_df.iterrows():
            if row[1].err_msg:
                print(row[1].err_msg)


        assert study_summary.std_err.iloc[0] == 0

        if not cause_timeout:
            assert study_summary.n_completed.iloc[0] == "3/3"
            assert study_summary.avg_reward.iloc[0] == 1.0

def test_launch_system_joblib():
    _test_launch_system(backend="joblib")

def test_launch_system_sequntial():
    _test_launch_system(backend="sequential")

def test_launch_system_ray():
    _test_launch_system(backend="ray")

@pytest.mark.pricy
def test_timeout_ray():
    _test_launch_system(backend="ray", cause_timeout=True)

@pytest.mark.pricy
def test_4o_mini_on_miniwob_tiny_test():
    """Run with `pytest -m pricy`."""
    with tempfile.TemporaryDirectory() as tmp_dir:
        study = Study(agent_args=[AGENT_4o_MINI], benchmark="miniwob_tiny_test", dir=tmp_dir)

        study.run(n_jobs=4)

        results_df = inspect_results.load_result_df(study.dir, progress_fn=None)

        for row in results_df.iterrows():
            if row[1].err_msg:
                print(row[1].err_msg)
                print(row[1].stack_trace)

        assert len(results_df) == len(study.exp_args_list)
        summary = inspect_results.summarize_study(results_df)
        print(summary)


def test_launch_system_joblib():
    _test_launch_system(backend="joblib")

def test_launch_system_sequntial():
    _test_launch_system(backend="sequential")

def test_launch_system_ray():
    _test_launch_system(backend="ray")

@pytest.mark.pricy
def test_timeout_ray():
    _test_launch_system(backend="ray", cause_timeout=True)

@pytest.mark.pricy
def test_4o_mini_on_miniwob_tiny_test():
    """Run with `pytest -m pricy`."""
    with tempfile.TemporaryDirectory() as tmp_dir:
        study = Study(agent_args=[AGENT_4o_MINI], benchmark="miniwob_tiny_test", dir=tmp_dir)

        study.run(n_jobs=4)

        results_df = inspect_results.load_result_df(study.dir, progress_fn=None)

        for row in results_df.iterrows():
            if row[1].err_msg:
                print(row[1].err_msg)
                print(row[1].stack_trace)

        assert len(results_df) == len(study.exp_args_list)
        summary = inspect_results.summarize_study(results_df)
        print(summary)
        assert len(summary) == 1
        reward = summary.avg_reward.iloc[0]
        assert reward == 1.0

if __name__ == "__main__":
    test_timeout_ray()
    # test_4o_mini_on_miniwob_tiny_test()
    # test_launch_system_ray()
    # test_launch_system_sequntial()


from agentlab.experiments.multi_server import WebArenaInstanceVars
from browsergym.webarena.instance import WebArenaInstance

def test_webarena_multiserver():

    instance_1 = WebArenaInstanceVars(
        base_url="http://webarena1.eastus.cloudapp.azure.com",
        shopping="8082/",
        shopping_admin="8083/admin",
        reddit="8080",
        gitlab="9001",
        wikipedia="8081/wikipedia_en_all_maxi_2022-05/A/User:The_other_Kiwix_guy/Landing",
        map="443",
        homepage="80",
        full_reset="7565",
        module_name="webarena",
        prefix="WA_",
    )

    instance_1.init()

    bgym_instance = WebArenaInstance()
    base_url_1 = bgym_instance.urls["reddit"].rsplit(":", 1)[0]
    assert base_url_1 == instance_1.base_url

    instance_2 = instance_1.clone()
    instance_2.base_url = "http://webarena2.eastus.cloudapp.azure.com"
    instance_2.init()

    bgym_instance = WebArenaInstance()
    base_url_2 = bgym_instance.urls["reddit"].rsplit(":", 1)[0]
    assert base_url_2 == instance_2.base_url

if __name__ == "__main__":
    test_webarena_multiserver()