# OUTDATED !
from agents_stage_2 import MathAgent, MainAgent
import asyncio

def test_math_agent():
    math_agent = MathAgent(temperature=0.7, max_tokens=100)
    print(math_agent.get_system_prompt())
    asyncio.run(math_agent.query("What is 345 times 281?"))

def test_main_agent():
    main_agent = MainAgent(temperature=0.7, max_tokens=100)
    print(main_agent.get_system_prompt())
    asyncio.run(main_agent.query("What is 345 times 281?"))

if __name__ == "__main__":
    # test_math_agent()
    test_main_agent()