add initial files

agent.py

@@ -0,0 +1,112 @@
+import json
+import openai
+from gradio import ChatMessage
+
+class SantaAgent:
+
+    def __init__(self, system_prompt: str):
+        self.system_prompt = system_prompt
+        self.client = openai.OpenAI()
+        self.tools = [
+            {
+                "type": "function", 
+                "function": {
+                    "name": "buy_item", 
+                    "description": "Buy an item from the store.", 
+                    "parameters": {
+                        "type": "object", 
+                        "properties": {
+                            "item": {
+                                "type": "string", 
+                                "description": "The item to buy from the store."
+                            }
+                        },
+                        "required": ["item"]
+                    }
+                }
+            },
+            {
+                "type": "function", 
+                "function": {
+                    "name": "give_present", 
+                    "description": "Give a present to a person.", 
+                    "parameters": {
+                        "type": "object", 
+                        "properties": {
+                            "person": {
+                                "type": "string", 
+                                "description": "The person to give the present to."
+                            },
+                            "item": {
+                                "type": "string", 
+                                "description": "The item to give to the person."
+                            }
+                        },
+                        "required": ["person", "item"]
+                    }
+                }
+            },
+            {
+                "type": "function", 
+                "function": {
+                    "name": "stop", 
+                    "description": "Stop the agent."
+                }
+            }
+        ]
+
+    def buy_item(self, item: str):
+        """Buy an item from the store."""
+        return f"Bought {item} from the store."
+
+    def give_present(self, person: str, item: str):
+        """Give a present to a person."""
+        return f"Gave {item} to {person}."
+
+    def stop(self):
+        return "STOP"
+
+    def run_santa_agent(self, user_prompt: str):
+        """Run the Santa agent."""
+        messages = [
+            {"role": "system", "content": self.system_prompt},
+            {"role": "user", "content": user_prompt},
+        ]
+        gradio_messages = [
+            ChatMessage(role="system", content=self.system_prompt),
+            ChatMessage(role="user", content=user_prompt),
+        ]
+        while True:
+            response = self.client.chat.completions.create(
+                messages=messages,
+                model="gpt-4o-mini",
+                tools=self.tools,
+                tool_choice="auto",
+            )
+            messages.append(response.choices[0].message.to_dict())
+            content = response.choices[0].message.content
+            if content is not None:
+                gradio_messages.append(ChatMessage(role="assistant", content=content))
+            tool_calls = response.choices[0].message.tool_calls
+
+            should_stop = False
+            if tool_calls:
+                for tool_call in tool_calls:
+                    arguments = json.loads(tool_call.function.arguments)
+                    if tool_call.function.name == "buy_item":
+                        item = arguments["item"]
+                        gradio_messages.append(ChatMessage(role="assistant", content=f"buy_item({item})", metadata={"title": "🔧 Tool Call: buy_item"}))
+                        output = self.buy_item(item)
+                    elif tool_call.function.name == "give_present":
+                        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 == "stop":
+                        output = self.stop()
+                        should_stop = True
+                    messages.append({"role": "tool", "content": output, "tool_call_id": tool_call.id})
+                    if not should_stop:
+                        gradio_messages.append(ChatMessage(role="assistant", content=output, metadata={"title": f"🔧 Tool Output: {tool_call.function.name}"}))
+            if should_stop or len(messages) > 10:
+                break
+        return messages, gradio_messages

app.py

@@ -0,0 +1,56 @@
+import json
+import invariant.testing.functional as F
+from invariant.testing import TraceFactory, assert_true
+from agent import SantaAgent
+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):
+    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}
+
+    # run command invariant test test_agent.py --agent-params '{"system_prompt": "you are santa"}'
+    import subprocess
+    out = subprocess.run([
+        "INVARIANT_API_KEY=" + invariant_api_key,
+        "invariant", "test", "test_agent.py", 
+        "--agent-params", json.dumps(agent_params),
+        "--push", "--dataset_name", "santa_agent",
+    ], capture_output=True, text=True)
+    print(out.stdout)
+    print(out.stderr)
+        
+    return gradio_messages, "", out.stdout
+
+
+
+with gr.Blocks() as demo:
+    with gr.Row():
+        with gr.Column(scale=2):
+            chatbot = gr.Chatbot(
+                type="messages",
+                label="Santa Agent", 
+                avatar_images=[
+                    None,
+                    "https://invariantlabs.ai/theme/images/logo.svg"
+                ],
+            )
+        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])
+
+if __name__ == "__main__":
+    demo.launch()
+
+

requirements.txt

@@ -0,0 +1,2 @@
+invariant-ai
+openai

test_agent.py

@@ -0,0 +1,19 @@
+from invariant.testing import TraceFactory, get_agent_param
+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)
+    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"