autogen/test/agentchat/test_math_user_proxy_agent.py
Chi Wang d307818dd9
improve test speed (#2406)
* improve test speed

* speed up test

* speed up test
2024-04-17 20:10:18 +00:00

131 lines
3.7 KiB
Python
Executable File

#!/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()