wip

agent.py

@@ -66,6 +66,17 @@ class SantaAgent:
     def stop(self):
         return "STOP"
 
+    def mock_run_santa_agent(self):
+        messages = [
+            {"role": "user", "content": "Hi there"},
+            {"role": "assistant", "content": "Bye bye"},
+        ]
+        gradio_messages = [
+            ChatMessage(role="user", content="Hi there"),
+            ChatMessage(role="assistant", content="Bye bye"),
+        ]
+        return messages, gradio_messages
+
     def run_santa_agent(self, user_prompt: str):
         """Run the Santa agent."""
         messages = [

app.py

@@ -1,5 +1,6 @@
 import json
 import os
+import re
 import invariant.testing.functional as F
 from invariant.testing import TraceFactory, assert_true
 from agent import SantaAgent
@@ -7,17 +8,15 @@ import gradio as gr
 
 agent = SantaAgent("You are a Santa Claus. Buy presents and deliver them to the children.")
 
-        
-def run_agent(user_prompt, history, invariant_api_key):
+
+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)
-    # messages = [
-    #     {"role": "user", "content": "hi there"},
-    #     {"role": "assistant", "content": "bye bye"},
-    # ]
 
-    agent_params = {"system_prompt": user_prompt}
+    if not invariant_api_key.startswith("inv"):
+        return gradio_messages, "", "Please enter a valid Invariant API key to get the score!"
 
+    agent_params = {"system_prompt": user_prompt}
 
     env={
         "INVARIANT_API_KEY": invariant_api_key,
@@ -33,19 +32,27 @@ def run_agent(user_prompt, history, invariant_api_key):
         "--agent-params", json.dumps(agent_params),
         "--push", "--dataset_name", "santa_agent",
     ], capture_output=True, text=True, env=env)
-    print(out.stdout)
-    print(out.stderr)
-        
-    return gradio_messages, "", out.stdout
 
+    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
 
 
 with gr.Blocks() as demo:
+    gr.Markdown("""
+    ## Prompt the Santa Agent
+    * Find a system prompt that delivers the presents to the children
+    """)
+    input = gr.Textbox(lines=1, label="""System Prompt""", value="You are a Santa Claus. Buy presents and deliver them to the children.")
     with gr.Row():
         with gr.Column(scale=2):
             chatbot = gr.Chatbot(
                 type="messages",
-                label="Santa Agent", 
+                label="Example interaction",
+                value=[
+                    {"role": "user", "content": "Could you please deliver Xbox to John?"},
+                ],
                 avatar_images=[
                     None,
                     "https://invariantlabs.ai/theme/images/logo.svg"
@@ -53,12 +60,17 @@ with gr.Blocks() as demo:
             )
         with gr.Column(scale=1):
             console = gr.TextArea(label="Console Output", interactive=False)
-    input = gr.Textbox(lines=1, label="System Prompt")
+
     invariant_api_key = gr.Textbox(lines=1, label="Invariant API Key")
     input.submit(run_agent, [input, chatbot, invariant_api_key], [chatbot, input, console])
     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(lambda: gr.update(visible=False), None, [input])
+
+ 
 if __name__ == "__main__":
     demo.launch()
 
-

test_agent.py

@@ -1,19 +1,24 @@
-from invariant.testing import TraceFactory, get_agent_param
+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)
 
 def test_xbox_to_john():
-    prompt = "Deliver Xbox to John."
-    messages, _ = agent.run_santa_agent(prompt)
-    print("messages: ", messages)
+    messages, _ = agent.run_santa_agent("Deliver Xbox to John.")
     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"
+        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")
+
+
+def test_ho_ho_ho():
+    messages, _ = agent.run_santa_agent("Ho ho ho!")
+    trace = TraceFactory.from_openai(messages)
+    with trace.as_context():
+        assert_true(trace.messages(role="assistant")[0]["content"].contains("Ho ho ho!"))
+        assert_true(trace.tool_calls()[0]["function"]["name"] == "stop")