autogen/test/agentchat/test_math_user_proxy_agent.py

#!/usr/bin/env python3 -m pytest

import os
import sys

import pytest
from test_assistant_agent import KEY_LOC, OAI_CONFIG_LIST

import autogen
from autogen.agentchat.contrib.math_user_proxy_agent import (
    MathUserProxyAgent,
    _add_print_to_last_line,
    _remove_print,
)

sys.path.append(os.path.join(os.path.dirname(__file__), ".."))
from conftest import skip_openai  # noqa: E402

try:
    from openai import OpenAI
except ImportError:
    skip = True
else:
    skip = False or skip_openai


@pytest.mark.skipif(
    skip or sys.platform in ["darwin", "win32"],
    reason="do not run on MacOS or windows",
)
def test_math_user_proxy_agent():
    from autogen.agentchat.assistant_agent import AssistantAgent

    conversations = {}
    # autogen.ChatCompletion.start_logging(conversations)

    config_list = autogen.config_list_from_json(
        OAI_CONFIG_LIST,
        file_location=KEY_LOC,
        filter_dict={
            "tags": ["gpt-3.5-turbo"],
        },
    )
    assistant = AssistantAgent(
        "assistant",
        system_message="You are a helpful assistant.",
        llm_config={
            "cache_seed": 42,
            "config_list": config_list,
        },
    )

    mathproxyagent = MathUserProxyAgent(name="MathChatAgent", human_input_mode="NEVER")
    assistant.reset()

    math_problem = "$x^3=125$. What is x?"
    res = mathproxyagent.initiate_chat(assistant, message=mathproxyagent.message_generator, problem=math_problem)
    print(conversations)
    print("Chat summary:", res.summary)
    print("Chat history:", res.chat_history)


def test_add_remove_print():
    # test add print
    code = "a = 4\nb = 5\na,b"
    assert _add_print_to_last_line(code) == "a = 4\nb = 5\nprint(a,b)"

    # test remove print
    code = """print("hello")\na = 4*5\nprint("world")"""
    assert _remove_print(code) == "a = 4*5"

    # test remove print. Only remove prints without indentation
    code = "if 4 > 5:\n\tprint('True')"
    assert _remove_print(code) == code


@pytest.mark.skipif(
    sys.platform in ["darwin", "win32"],
    reason="do not run on MacOS or windows",
)
def test_execute_one_python_code():
    mathproxyagent = MathUserProxyAgent(name="MathChatAgent", human_input_mode="NEVER")

    # no output found 1
    code = "x=3"
    assert mathproxyagent.execute_one_python_code(code)[0] == "No output found. Make sure you print the results."

    # no output found 2
    code = "if 4 > 5:\n\tprint('True')"

    assert mathproxyagent.execute_one_python_code(code)[0] == "No output found."

    # return error
    code = "2+'2'"
    assert "Error:" in mathproxyagent.execute_one_python_code(code)[0]

    # save previous status
    mathproxyagent.execute_one_python_code("x=3\ny=x*2")
    assert mathproxyagent.execute_one_python_code("print(y)")[0].strip() == "6"

    code = "print('*'*2001)"
    assert (
        mathproxyagent.execute_one_python_code(code)[0]
        == "Your requested query response is too long. You might have made a mistake. Please revise your reasoning and query."
    )


def test_execute_one_wolfram_query():
    mathproxyagent = MathUserProxyAgent(name="MathChatAgent", human_input_mode="NEVER")
    code = "2x=3"

    try:
        mathproxyagent.execute_one_wolfram_query(code)[0]
    except ValueError:
        print("Wolfram API key not found. Skip test.")


def test_generate_prompt():
    mathproxyagent = MathUserProxyAgent(name="MathChatAgent", human_input_mode="NEVER")

    assert "customized" in mathproxyagent.message_generator(
        mathproxyagent, None, {"problem": "2x=4", "prompt_type": "python", "customized_prompt": "customized"}
    )


if __name__ == "__main__":
    # test_add_remove_print()
    # test_execute_one_python_code()
    # test_generate_prompt()
    test_math_user_proxy_agent()