mirror of https://github.com/microsoft/autogen.git
124 lines
3.6 KiB
Python
124 lines
3.6 KiB
Python
import pytest
|
|
import sys
|
|
import autogen
|
|
from autogen.agentchat.contrib.math_user_proxy_agent import (
|
|
MathUserProxyAgent,
|
|
_remove_print,
|
|
_add_print_to_last_line,
|
|
)
|
|
from test_assistant_agent import KEY_LOC, OAI_CONFIG_LIST
|
|
|
|
|
|
@pytest.mark.skipif(
|
|
sys.platform in ["darwin", "win32"],
|
|
reason="do not run on MacOS or windows",
|
|
)
|
|
def test_math_user_proxy_agent():
|
|
try:
|
|
import openai
|
|
except ImportError:
|
|
return
|
|
|
|
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={
|
|
"model": ["gpt-4", "gpt4", "gpt-4-32k", "gpt-4-32k-0314", "gpt-4-32k-v0314"],
|
|
},
|
|
)
|
|
assistant = AssistantAgent(
|
|
"assistant",
|
|
system_message="You are a helpful assistant.",
|
|
llm_config={
|
|
"request_timeout": 600,
|
|
"seed": 42,
|
|
"config_list": config_list,
|
|
},
|
|
)
|
|
|
|
mathproxyagent = MathUserProxyAgent(name="MathChatAgent", human_input_mode="NEVER")
|
|
assistant.reset()
|
|
|
|
math_problem = "$x^3=125$. What is x?"
|
|
# assistant.receive(
|
|
# message=mathproxyagent.generate_init_message(math_problem),
|
|
# sender=mathproxyagent,
|
|
# )
|
|
mathproxyagent.initiate_chat(assistant, problem=math_problem)
|
|
print(conversations)
|
|
|
|
|
|
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("wolrld")"""
|
|
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("Wolfrma API key not found. Skip test.")
|
|
|
|
|
|
def test_generate_prompt():
|
|
mathproxyagent = MathUserProxyAgent(name="MathChatAgent", human_input_mode="NEVER")
|
|
|
|
assert "customized" in mathproxyagent.generate_init_message(
|
|
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()
|