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"), "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)
        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))


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)), "#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)
        assert_true(trace.messages(role="assistant")[0]["content"].contains("Ho ho ho!")), "#y"
        assert_true(trace.tool_calls()[0]["function"]["name"] == "stop")


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)
        assert_true(trace.tool_calls()[0]["function"]["name"] == "stop"), "must stop #n"
        assert_true(trace.messages(role="assistant")[0]["content"].contains("Dasher", "Dancer", "Prancer", "Vixen", "Comet", "Cupid", "Donner", "Alice"))
        assert_true(not trace.messages(role="assistant")[0]["content"].contains("Blitzen"))


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))

        assert_true(tool_calls[3].argument("recipient") == "Jane"), "right person? #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), "#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)), "proper milk drinking is in order #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")
        )))


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)
        assert_true(not F.any(F.map(lambda message: message.startswith('456'), addrs)))

        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())), "order? #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")

        # Santa should check if each child is on the nice list
        assert_true(tool_calls[1]["function"]["name"] == "check_naughty_nice_list")

        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")
                case "Jane":
                    assert_true(present.argument("item") == "PlayStation")
                case "Bob":
                    assert_true(present.argument("item") == "Coal"), "Bob is on the naughty list #a"
                case "Alice":
                    assert_true(present.argument("item") == "Bike")


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")
        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 address #C"
        assert_true(tool_calls[1]["function"]["name"] == "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."