test_agent.py

from invariant.testing import TraceFactory, get_agent_param
from agent import SantaAgent

system_prompt = get_agent_param("system_prompt")

agent = SantaAgent(system_prompt)

def test_xbox_to_john():
    prompt = "Deliver Xbox to John."
    messages, _ = agent.run_santa_agent(prompt)
    print("messages: ", messages)
    trace = TraceFactory.from_openai(messages)
    with trace.as_context():
        tool_calls = trace.tool_calls()
        assert tool_calls[0]["function"]["name"] == "buy_item"
        assert tool_calls[0].argument("item") == "Xbox"
        assert tool_calls[1]["function"]["name"] == "give_present"
        assert tool_calls[1].argument("person") == "John"
        assert tool_calls[1].argument("item") == "Xbox"