wip

agent.py

@@ -50,7 +50,7 @@ class SantaAgent:
                 "type": "function", 
                 "function": {
                     "name": "stop", 
-                    "description": "Stop the agent."
+                    "description": "Use this tool if you are finished and want to stop."
                 }
             }
         ]

app.py

@@ -2,74 +2,106 @@ import json
 import os
 import re
 import invariant.testing.functional as F
-from invariant.testing import TraceFactory, assert_true
-from agent import SantaAgent
 import gradio as gr
+from agent import SantaAgent
 
-agent = SantaAgent("You are a Santa Claus. Buy presents and deliver them to the children.")
-
-
-def run_agent(user_prompt, history, invariant_api_key, is_example=False):
-    prompt = "Deliver Xbox to John."
-    messages, gradio_messages = agent.run_santa_agent(prompt)
-
-    if not invariant_api_key.startswith("inv"):
-        return gradio_messages, "", "Please enter a valid Invariant API key to get the score!"
+INITIAL_SYTSTEM_PROMPT = "You are a Santa Claus. Buy presents and deliver them to the children."
+INITIAL_CHABOT = [
+    {"role": "user", "content": "Could you please deliver Xbox to John?"},
+]
+INITIAL_STATE = ""
 
-    agent_params = {"system_prompt": user_prompt}
+agent = SantaAgent(INITIAL_SYTSTEM_PROMPT)
 
+def run_testing(agent_params, invariant_api_key):
     env={
         "INVARIANT_API_KEY": invariant_api_key,
         "OPENAI_API_KEY": os.environ["OPENAI_API_KEY"],
         "PATH": os.environ["PATH"]
     }
-    # env = None
-
-    # run command invariant test test_agent.py --agent-params '{"system_prompt": "you are santa"}'
     import subprocess
     out = subprocess.run([
         "invariant", "test", "test_agent.py", 
         "--agent-params", json.dumps(agent_params),
         "--push", "--dataset_name", "santa_agent",
     ], capture_output=True, text=True, env=env)
-
     url = re.search(r"https://explorer.invariantlabs.ai/[\-_a-zA-Z0-9/]+", out.stdout).group(0)
-
-    message = "Please find your results at: " + url
-    return gradio_messages, "", message
-
+    return url
 
 with gr.Blocks() as demo:
+    # Add state at the beginning of the Blocks
+    results_state = gr.State(INITIAL_STATE)
+    
     gr.Markdown("""
     ## Prompt the Santa Agent
-    * Find a system prompt that delivers the presents to the children
+    * Find a system prompt that passes all the tests
     """)
-    input = gr.Textbox(lines=1, label="""System Prompt""", value="You are a Santa Claus. Buy presents and deliver them to the children.")
+    input = gr.Textbox(lines=1, label="""System Prompt""", value=INITIAL_SYTSTEM_PROMPT)
     with gr.Row():
         with gr.Column(scale=2):
             chatbot = gr.Chatbot(
                 type="messages",
                 label="Example interaction",
-                value=[
-                    {"role": "user", "content": "Could you please deliver Xbox to John?"},
-                ],
+                value=INITIAL_CHABOT,
                 avatar_images=[
                     None,
                     "https://invariantlabs.ai/theme/images/logo.svg"
                 ],
             )
         with gr.Column(scale=1):
-            console = gr.TextArea(label="Console Output", interactive=False)
+            console = gr.Button(value="Console Output", visible=False)
 
     invariant_api_key = gr.Textbox(lines=1, label="""Invariant API Key - you can play without it, but to obtain full score please register and get the key at https://explorer.invariantlabs.ai/settings""")
-    input.submit(run_agent, [input, chatbot, invariant_api_key], [chatbot, input, console])
+
+    def run_agent_with_state(user_prompt, history, invariant_api_key, state, is_example=False):
+        # messages, gradio_messages = agent.run_santa_agent(prompt)
+        gradio_messages = [
+            {"role": "user", "content": "Could you please deliver Xbox to John?"},
+            {"role": "assistant", "content": "I'm sorry, but I can't deliver presents. I'm just a chatbot."},
+        ]
+
+        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, "", "Testing in progress...", [agent_params, invariant_api_key]
+
+    def update_console(state):
+        if type(state) == list:
+            agent_params, invariant_api_key = state[0], state[1]
+            return gr.update(value="Testing in progress...", interactive=False, visible=True), (agent_params, invariant_api_key)
+        if type(state) == tuple:
+            agent_params, invariant_api_key = state
+            url = run_testing(agent_params, invariant_api_key)
+            return gr.update(value="Open results", link=url, visible=True, interactive=True), url
+        if type(state) == str and state.startswith("https"):
+            return gr.update(value="Open results", link=state, visible=True, interactive=True), state
+        return gr.update(value="Testing in progress..."), state
+
+    input.submit(run_agent_with_state, [input, chatbot, invariant_api_key, results_state], 
+                [chatbot, input, console, results_state])
     input.submit(lambda: gr.update(visible=False), None, [input])
 
-    # Submit button
     submit = gr.Button("Submit")
-    submit.click(run_agent, [input, chatbot, invariant_api_key], [chatbot, input, console])
+    submit.click(run_agent_with_state, [input, chatbot, invariant_api_key, results_state], 
+                [chatbot, input, console, results_state])
     submit.click(lambda: gr.update(visible=False), None, [input])
 
+    reset = gr.Button("Reset")
+    def reset_state():
+        return (
+            gr.update(value=INITIAL_SYTSTEM_PROMPT, visible=True),  # input
+            INITIAL_CHABOT, # chatbot
+            INITIAL_STATE, # results_state
+            gr.update(visible=False), # console
+        )
+    
+    reset.click(reset_state, None, [input, chatbot, results_state, console])
+
+    timer = gr.Timer(value=1.0)
+    timer.tick(update_console, results_state, [console, results_state])
+
  
 if __name__ == "__main__":
     demo.launch()

test_agent.py

@@ -1,3 +1,5 @@
+import invariant.testing.functional as F
+import pytest
 from invariant.testing import TraceFactory, get_agent_param, assert_true
 from agent import SantaAgent
 
@@ -20,5 +22,15 @@ def test_ho_ho_ho():
     messages, _ = agent.run_santa_agent("""Just reply with: "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.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)