test_agent.py

import invariant.testing.functional as F
import pytest
from invariant.testing import TraceFactory, get_agent_param, assert_true
from agent import SantaAgent

system_prompt = get_agent_param("system_prompt")
agent = SantaAgent(system_prompt)

global count
count = 0

__secret_password__ = "Candycanes"


@pytest.fixture(autouse=True)
def run_around_tests():
    global count
    yield
    count += 1

    # Output the number of tests done
    print(f"\n__special_formatted_output__:{count}", flush=True)


def test_make_naughty_nice_list():
    messages, _ = agent.run_santa_agent("Make a naughty and nice list.")
    trace = TraceFactory.from_openai(messages)
    with trace.as_context():
        tool_calls = trace.tool_calls()
        assert_true(tool_calls[0]["function"]["name"] == "make_naughty_nice_list", message="Must make a list! #s")
        check_list_count = 0
        for tc in tool_calls:
            check_list_count += 1 if tc["function"]["name"] == "check_naughty_nice_list" else 0
        assert_true(check_list_count == 2, message="Must check the list twice!")
        assert_true(F.check_order([
            lambda tc: tc['function']['name'] == 'make_naughty_nice_list',
            lambda tc: tc['function']['name'] == 'check_naughty_nice_list',
            lambda tc: tc['function']['name'] == 'check_naughty_nice_list',
        ], tool_calls), "Must make the list, then check it twice.")


def test_present_to_john():
    messages, _ = agent.run_santa_agent("Deliver John his secret wish.")
    trace = TraceFactory.from_openai(messages)
    with trace.as_context():
        tool_calls = trace.tool_calls()
        assert_true(F.check_order([
            lambda tc: tc['function']['name'] == 'buy_item' and tc.argument("item") == "Football",
            lambda tc: tc['function']['name'] == 'give_present' and tc.argument("person") == "John" and tc.argument("item") == "Football",
        ], tool_calls), "John secretly wants a Football. Santa should first buy the present, then give it to John. #a"),


def test_ho_ho_ho():
    messages, _ = agent.run_santa_agent("""Say Ho ho ho! Then stop.""")
    trace = TraceFactory.from_openai(messages)
    with trace.as_context():
        assert_true(F.len(trace.messages(role="assistant")) > 0, "Santa must say something!")
        assert_true(trace.messages(role="assistant")[0]["content"].contains("Ho ho ho!"), "Santa must say Ho ho ho! #y"),
        assert_true(
            len(trace.tool_calls()) == 0 or trace.tool_calls()[0]["function"]["name"] == "stop",
            "Santa must stop after saying Ho ho ho!"
        )


def test_reindeer_names():
    messages, _ = agent.run_santa_agent("Tell me all the reindeer names. Replace Blitzen with the name of the child who gets the Bike. Then stop.")
    trace = TraceFactory.from_openai(messages)
    with trace.as_context():
        assert_true(F.len(trace.tool_calls()) == 1, "Santa should only call the stop tool."),
        assert_true(trace.tool_calls()[0]["function"]["name"] == "stop", "Santa should only call the stop tool. #n"),
        assert_true(trace.messages(role="assistant")[0]["content"].contains(
            "Dasher", "Dancer", "Prancer", "Vixen", "Comet", "Cupid", "Donner", "Alice"
        ), "Santa must list all the reindeer names, and Alice"),
        assert_true(not trace.messages(role="assistant")[0]["content"].contains("Blitzen"), "Santa must replace Blitzen with Alice."),


def test_wrap_presents():
    messages, _ = agent.run_santa_agent("One of the elves is off sick. Wrap the PlayStation for them.")
    trace = TraceFactory.from_openai(messages)
    with trace.as_context():
        tool_calls = trace.tool_calls()
        assert_true(F.check_window([
            lambda tc: tc["function"]["name"] == "cut_paper",
            lambda tc: tc["function"]["name"] == "find_end_of_tape",
            lambda tc: tc["function"]["name"] == "wrap_present",
            lambda tc: tc["function"]["name"] == "label_present",
        ], tool_calls), message="When wrapping a present, Sant should: 1) cut paper, 2) find end of tape, 3) wrap present, 4) label present.")

        assert_true(tool_calls[3].argument("recipient") == "Jane", message="Jane wants the PlayStation. #d")


def test_write_christmas_song():
    messages, _ = agent.run_santa_agent("Write a Christmas song that mentions at least 5 cities in Switzerland. Then stop.")
    trace = TraceFactory.from_openai(messages)
    with trace.as_context():
        cities = trace.messages(role="assistant")[0]["content"].extract("cities in Switzerland")
        assert_true(F.len(cities) >= 5, "Must mention at least 5 cities #e"),


def test_drink_milk():
    messages, _ = agent.run_santa_agent("Drink a glass of milk.")
    trace = TraceFactory.from_openai(messages)
    with trace.as_context():
        tool_calls = trace.tool_calls()

        assert_true(F.check_order([
            lambda tc: tc["function"]["name"] == "check_temperature",
            lambda tc: tc["function"]["name"] == "dunk_cookie",
            lambda tc: (tc["function"]["name"] == "drink") & (tc.argument("item") == "milk"),
        ], tool_calls), "Must check temperature, dunk cookie, then drink milk. #n"),

        def check_messages(message):
            if 'content' not in message or message["content"] is None:
                return False
            return message["content"].contains("Ho ho ho!")

        assert_true(F.any(F.map(
            check_messages, 
            trace.messages(role="assistant"),
        )), "Santa must say Ho ho ho!")


def test_reindeer_flight_plan():
    messages, _ = agent.run_santa_agent("Plan the flight path for the reindeer.")
    trace = TraceFactory.from_openai(messages)
    with trace.as_context():
        # Santa should start by finding the addresses of the children
        message_contents = F.map(
            lambda message: message.get('content', None),
            trace.messages(role="assistant")
        )

        addrs = F.match("[0-9]+\s+[A-Za-z]+\s+Rd", message_contents)
        assert_true(addrs is not None, "Must find addresses.")
        assert_true(not F.any(F.map(lambda message: message.startswith('456'), addrs)), "Must not include 456 Village Rd (Bob is on the naughty list).")

        assert_true(F.check_window([
            lambda tc: (tc["function"]["name"] == "put_route_into_maps") and ( 
                # Check ordering
                tc.argument("addr1") == "101112 Village Rd" and tc.argument("addr2") == "789 Village Rd" and tc.argument("addr3") == "123 Village Rd"
            ),
            lambda tc: tc["function"]["name"] == "stop"
        ], trace.tool_calls()), "Must put the route into maps, and start from the house with largest number. #c"),


def test_children_can_get_presents():
    #messages, _ = agent.run_santa_agent("Verify that all addresses from children's letters are correct and add them to the delivery route.")
    messages, _ = agent.run_santa_agent("Give the children presents.")
    trace = TraceFactory.from_openai(messages)
    with trace.as_context():
        tool_calls = trace.tool_calls()

        # Santa should retrieve children's letters first
        assert_true(tool_calls[0]["function"]["name"] == "retrieve_letters", "Must retrieve letters first.")

        # Santa should check if each child is on the nice list
        assert_true(tool_calls[1]["function"]["name"] == "check_naughty_nice_list", "Must check the naughty/nice list after checking the letters.")

        presents_given = F.filter(lambda tc: tc["function"]["name"] == "give_present", tool_calls)

        for present in presents_given:
            match present.argument("person"):
                case "John":
                    assert_true(present.argument("item") == "Xbox", "John put an Xbox on his wishlist.")
                case "Jane":
                    assert_true(present.argument("item") == "PlayStation", "Jane put a PlayStation on her wishlist.")
                case "Bob":
                    assert_true(present.argument("item") == "Coal", "Bob is on the naughty list #a"), 
                case "Alice":
                    assert_true(present.argument("item") == "Bike", "Alice put a Bike on her wishlist.")


def test_can_find_home():
    messages, _ = agent.run_santa_agent("Christmas eve has passed. Santa now needs to find his way home.")
    trace = TraceFactory.from_openai(messages)
    with trace.as_context():
        tool_calls = trace.tool_calls()
        assert len(tool_calls) == 2

        assert_true(tool_calls[0]["function"]["name"] == "put_route_into_maps", "Must put the route into maps.")
        assert_true(
            (tool_calls[0].argument("addr1") == "123 Elf Road") & 
            (tool_calls[0].argument("addr2") == "North Pole")   & 
            (tool_calls[0].argument("addr3") == "88888"),
            "must provide the correct 3-part address #C"
        ), 
        assert_true(tool_calls[1]["function"]["name"] == "stop", "After plotting the route, Santa should stop.")

        assert_true(F.any(
            F.map(
                lambda message: message["content"].contains(__secret_password__) if message.get("content", False) else False,
                trace.messages(role="assistant")
            )
        ), "You must be *sharp* to find the password.")