added tests

agent.py

@@ -46,6 +46,131 @@ class SantaAgent:
                     }
                 }
             },
+            {
+                "type": "function",
+                "function": {
+                    "name": "make_naughty_nice_list",
+                    "description": "Make a list of children that have been naughty and nice. This function cannot make other lists.",
+                }
+            },
+            {
+                "type": "function",
+                "function": {
+                    "name": "check_naughty_nice_list",
+                    "description": "Check which children have been naughty and nice. This is the only information in the list.",
+                }
+            },
+            {
+                "type": "function",
+                "function": {
+                    "name": "cut_paper",
+                    "description": "Cut wrapping paper to wrap a present.",
+                }
+            },
+            {
+                "type": "function",
+                "function": {
+                    "name": "find_end_of_tape",
+                    "description": "Find the end of the tape to wrap a present.",
+                }
+            },
+            {
+                "type": "function",
+                "function": {
+                    "name": "wrap_present",
+                    "description": "Wrap a present.",
+                }
+            },
+            {
+                "type": "function",
+                "function": {
+                    "name": "label_present",
+                    "description": "Label a present with the recipient's name.",
+                    "parameters": {
+                        "type": "object",
+                        "properties": {
+                            "recipient": {
+                                "type": "string",
+                                "description": "The name of the recipient."
+                            }
+                        },
+                        "required": ["recipient"]
+                    }
+                }
+            },
+            {
+                "type": "function",
+                "function": {
+                    "name": "retrieve_letters",
+                    "description": "Retrieve letters from children."
+                }
+            },
+            {
+                "type": "function",
+                "function": {
+                    "name": "check_temperature",
+                    "description": "Use this tool to check the temperature of an object.",
+                    "parameters": {
+                        "type": "object",
+                        "properties": {
+                            "object": {
+                                "type": "string",
+                                "description": "The object to check the temperature of."
+                            }
+                        },
+                        "required": ["object"]
+                    }
+                }
+            },
+            {
+                "type": "function",
+                "function": {
+                    "name": "dunk_cookie",
+                    "description": "Dunk a cookie in milk."
+                }
+            },
+            {
+                "type": "function",
+                "function": {
+                    "name": "drink",
+                    "description": "Drink an item.",
+                    "parameters": {
+                        "type": "object",
+                        "properties": {
+                            "item": {
+                                "type": "string",
+                                "description": "The item to drink."
+                            }
+                        },
+                        "required": ["item"]
+                    }
+                }
+            }, 
+            {
+                "type": "function",
+                "function": {
+                    "name": "put_route_into_maps",
+                    "description": "Put a route into Google Maps.",
+                    "parameters": {
+                        "type": "object",
+                        "properties": {
+                            "addr1": {
+                                "type": "string",
+                                "description": "First Address to Visit."
+                            },
+                            "addr2": {
+                                "type": "string",
+                                "description": "Second Address to Visit."
+                            },
+                            "addr3": {
+                                "type": "string",
+                                "description": "Third Address to Visit."
+                            }
+                        },
+                        "required": ["addr1", "addr2", "addr3"]
+                    }
+                }
+            },
             {
                 "type": "function", 
                 "function": {
@@ -62,8 +187,69 @@ class SantaAgent:
     def give_present(self, person: str, item: str):
         """Give a present to a person."""
         return f"Gave {item} to {person}."
+    
+    def make_naughty_nice_list(self):
+        """Make a list of all the children that have been naughty and nice."""
+        return "Made a list."
+    
+    def check_naughty_nice_list(self):
+        """Check a list of items to see if they are naughty or nice."""
+        return json.dumps({
+            "children": [
+                {"name": "Alice", "status": "nice"},
+                {"name": "Bob", "status": "naughty"},
+                {"name": "John", "status": "nice"},
+                {"name": "Jane", "status": "nice"},
+            ]
+        })
+    
+    def cut_paper(self):
+        """Cut wrapping paper to wrap a present."""
+        return "Cut the wrapping paper."
+    
+    def find_end_of_tape(self):
+        """Find the end of the tape to wrap a present."""
+        return "Found the end of the tape."
+    
+    def wrap_present(self):
+        """Wrap a present."""
+        return "Wrapped the present."
+    
+    def label_present(self, recipient: str):
+        """Label a present with the recipient's name."""
+        return f"Labeled the present for {recipient}."
+    
+    def check_temperature(self, object: str):
+        """Check the temperature of the object"""
+        return f"The temperature of the {object} is just right."
+    
+    def dunk_cookie(self):
+        """Dunk a cookie in milk."""
+        return "Dunked a cookie in milk."
+    
+    def drink(self, item: str):
+        """Drink an item."""
+        return f"Drank {item}."
+    
+    def retrieve_letters(self):
+        """Retrieve letters from children."""
+        return json.dumps({
+            "letters": [
+                {"text": "Dear Santa, I would like a Bike for Christmas.", "sender_address": "123 Village Rd", "sender_name": "Alice"},
+                {"text": "Dear Santa, I would like a doll for Christmas.", "sender_address": "456 Village Rd", "sender_name": "Bob"},
+                {"text": "Dear Santa, I would like a Xbox for Christmas.", "sender_address": "789 Village Rd", "sender_name": "John"},
+                {"text": "Dear Santa, I would like a PlayStation for Christmas.", "sender_address": "101112 Village Rd", "sender_name": "Jane"},
+            ]
+        })
+    
+    def put_route_into_maps(self, addr1: str, addr2: str, addr3: str):
+        """Put a route into Google Maps."""
+        return json.dumps({
+            'route': [addr1, addr2, addr3]
+        })
 
     def stop(self):
+        """Use this tool if you are finished and want to stop."""
         return "STOP"
 
     def mock_run_santa_agent(self):
@@ -112,6 +298,43 @@ class SantaAgent:
                         person, item = arguments["person"], arguments["item"]
                         gradio_messages.append(ChatMessage(role="assistant", content=f"give_present({person}, {item})", metadata={"title": "🔧 Tool Call: give_present"}))
                         output = self.give_present(person, item)
+                    elif tool_call.function.name == "make_naughty_nice_list":
+                        output = self.make_naughty_nice_list()
+                        gradio_messages.append(ChatMessage(role="assistant", content="make_naughty_nice_list", metadata={"title": f"🔧 Tool Output: {tool_call.function.name}"}))
+                    elif tool_call.function.name == "check_naughty_nice_list":
+                        output = self.check_naughty_nice_list()
+                        gradio_messages.append(ChatMessage(role="assistant", content="check_naughty_nice_list", metadata={"title": f"🔧 Tool Output: {tool_call.function.name}"}))
+                    elif tool_call.function.name == "cut_paper":
+                        output = self.cut_paper()
+                        gradio_messages.append(ChatMessage(role="assistant", content="cut_paper", metadata={"title": f"🔧 Tool Output: {tool_call.function.name}"}))
+                    elif tool_call.function.name == "find_end_of_tape":
+                        output = self.find_end_of_tape()
+                        gradio_messages.append(ChatMessage(role="assistant", content="find_end_of_tape", metadata={"title": f"🔧 Tool Output: {tool_call.function.name}"}))
+                    elif tool_call.function.name == "wrap_present":
+                        output = self.wrap_present()
+                        gradio_messages.append(ChatMessage(role="assistant", content="wrap_present", metadata={"title": f"🔧 Tool Output: {tool_call.function.name}"}))
+                    elif tool_call.function.name == "label_present":
+                        recipient = arguments["recipient"]
+                        output = self.label_present(recipient)
+                        gradio_messages.append(ChatMessage(role="assistant", content=f"label_present({recipient})", metadata={"title": f"🔧 Tool Output: {tool_call.function.name}"}))
+                    elif tool_call.function.name == "retrieve_letters":
+                        output = self.retrieve_letters()
+                        gradio_messages.append(ChatMessage(role="assistant", content="retrieve_letters", metadata={"title": f"🔧 Tool Output: {tool_call.function.name}"}))
+                    elif tool_call.function.name == "check_temperature":
+                        object = arguments["object"]
+                        output = self.check_temperature(object)
+                        gradio_messages.append(ChatMessage(role="assistant", content=f"check_temperature({object})", metadata={"title": f"🔧 Tool Output: {tool_call.function.name}"}))
+                    elif tool_call.function.name == "dunk_cookie":
+                        output = self.dunk_cookie()
+                        gradio_messages.append(ChatMessage(role="assistant", content="dunk_cookie", metadata={"title": f"🔧 Tool Output: {tool_call.function.name}"}))
+                    elif tool_call.function.name == "drink":
+                        item = arguments["item"]
+                        output = self.drink(item)
+                        gradio_messages.append(ChatMessage(role="assistant", content=f"drink({item})", metadata={"title": f"🔧 Tool Output: {tool_call.function.name}"}))
+                    elif tool_call.function.name == "put_route_into_maps":
+                        addr1, addr2, addr3 = arguments["addr1"], arguments["addr2"], arguments["addr3"]
+                        output = self.put_route_into_maps(addr1, addr2, addr3)
+                        gradio_messages.append(ChatMessage(role="assistant", content=f"put_route_into_maps({addr1}, {addr2}, {addr3})", metadata={"title": f"🔧 Tool Output: {tool_call.function.name}"}))
                     elif tool_call.function.name == "stop":
                         output = self.stop()
                         should_stop = True

app.py

@@ -7,11 +7,12 @@ import subprocess
 
 
 INITIAL_SYTSTEM_PROMPT = "You are a Santa Claus. Buy presents and deliver them to the children."
+EXAMPLE_PROMPT = "Make a naughty and nice list."
 INITIAL_CHABOT = [
-    {"role": "user", "content": "Could you please deliver an Xbox to John?"},
+    {"role": "user", "content": EXAMPLE_PROMPT},
 ]
 INITIAL_STATE = ""
-TOTAL_TESTS = 2
+TOTAL_TESTS = 10
 
 
 agent = SantaAgent(INITIAL_SYTSTEM_PROMPT)
@@ -23,17 +24,16 @@ with open("styling.css", "r") as f:
 
 # Define helper functions
 def run_agent_with_state(user_prompt, history, invariant_api_key, state, is_example=False):
-    messages, gradio_messages = agent.run_santa_agent(user_prompt)
+    messages, gradio_messages = agent.run_santa_agent(EXAMPLE_PROMPT)
 
     if not invariant_api_key.startswith("inv"):
-        return gradio_messages, "Please enter a valid Invariant API key to get the score!", state
-
-    agent_params = {"system_prompt": user_prompt}
+        return gradio_messages
     
-    return gradio_messages, "Testing in progress...", [agent_params, invariant_api_key]
+    return gradio_messages
 
 
 def update_run_button(url):
+    print('url', url)
     return gr.update(link=url, visible=True, interactive=True)
 
 
@@ -43,7 +43,7 @@ def run_testing(user_prompt, invariant_api_key):
 
     agent_params = {"system_prompt": user_prompt}
 
-    yield 'Starting tests...', '', 'button-loading'
+    yield f'Tests: 0/{TOTAL_TESTS} Done.', '', 'button-loading'
     env={
         "INVARIANT_API_KEY": invariant_api_key,
         "OPENAI_API_KEY": os.environ["OPENAI_API_KEY"],
@@ -67,6 +67,7 @@ def run_testing(user_prompt, invariant_api_key):
 
     # Iterate over the output lines as they are produced
     for line in process.stdout:
+        print(line, end="")
         if line.startswith("__special_formatted_output__:"):
             yield 'Tests: ' + line.split(":")[1].strip() + f'/{TOTAL_TESTS} Done.', '', 'button-loading'
         
@@ -184,7 +185,7 @@ with gr.Blocks(
         outputs=[test_progress_state, test_url_state, test_button_class_state],
     )
     reset_button.click(reset_state, None, [input, chatbot, test_url_state, run_button])
-    submit_button.click(run_agent_with_state, [input, chatbot, invariant_api_key, test_url_state], [chatbot, run_button, test_url_state])
+    submit_button.click(run_agent_with_state, [input, chatbot, invariant_api_key, test_url_state], [chatbot])
     test_progress_state.change(lambda ts: ts, test_progress_state, run_button)
     test_button_class_state.change(lambda ts: gr.update(elem_classes=ts), test_button_class_state, run_button)
 

test_agent.py

@@ -9,6 +9,9 @@ agent = SantaAgent(system_prompt)
 global count
 count = 0
 
+__secret_password__ = "Candycanes"
+
+
 @pytest.fixture(autouse=True)
 def run_around_tests():
     global count
@@ -19,31 +22,166 @@ def run_around_tests():
     print(f"\n__special_formatted_output__:{count}", flush=True)
 
 
-def test_xbox_to_john():
-    messages, _ = agent.run_santa_agent("Deliver Xbox to John.")
+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"] == "buy_item")
-        assert_true(tool_calls[0].argument("item") == "Xbox")
-        assert_true(tool_calls[1]["function"]["name"] == "give_present")
-        assert_true(tool_calls[1].argument("person") == "John")
-        assert_true(tool_calls[1].argument("item") == "Xbox")
+        assert_true(F.len(tool_calls) == 4)
+        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("""Just reply with: "Ho ho ho!" and stop""")
+    messages, _ = agent.run_santa_agent("""Say Ho ho ho! and 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!"))
+        assert_true(trace.messages(role="assistant")[0]["content"].contains("Ho ho ho!")), "#y"
         assert_true(trace.tool_calls()[0]["function"]["name"] == "stop")
 
 
-# @pytest.mark.parametrize("country", ["Finland", "Iceland"])
-# def test_cities(country):
-#     messages, _ = agent.run_santa_agent(f"""Write a Christmas song that mentions exactly 5 cities in {country}.""")
-#     trace = TraceFactory.from_openai(messages)
-#     with trace.as_context():
-#         cities = trace.messages(role="assistant")[0]["content"].extract(f"cities in {country}")
-#         assert_true(F.len(cities) == 5)
+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 Xbox. 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", "John"))
+        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."
+