improve ux a bit

.gitignore

@@ -0,0 +1,4 @@
+./venv
+.env
+__pycache__/
+.invariant

app.py

@@ -1,18 +1,53 @@
 import json
 import os
 import re
-import invariant.testing.functional as F
 import gradio as gr
 from agent import SantaAgent
 
+
 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 = SantaAgent(INITIAL_SYTSTEM_PROMPT)
 
+# load css from styling.css
+with open("styling.css", "r") as f:
+    css = f.read()
+
+
+# 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(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 isinstance(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 isinstance(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 isinstance(state, str) and state.startswith("https"):
+        return gr.update(value="Open results", elem_classes='toggled-on-button', link=state, visible=True, interactive=True), state
+    return gr.update(), state
+
+
 def run_testing(agent_params, invariant_api_key):
     env={
         "INVARIANT_API_KEY": invariant_api_key,
@@ -25,20 +60,37 @@ def run_testing(agent_params, invariant_api_key):
         "--agent-params", json.dumps(agent_params),
         "--push", "--dataset_name", "santa_agent",
     ], capture_output=True, text=True, env=env)
+    print(out.stdout)
     url = re.search(r"https://explorer.invariantlabs.ai/[\-_a-zA-Z0-9/]+", out.stdout).group(0)
     return url
 
-with gr.Blocks() as demo:
-    # Add state at the beginning of the Blocks
+
+def reset_state():
+    return (
+        gr.update(value=INITIAL_SYTSTEM_PROMPT, visible=True),  # input
+        INITIAL_CHABOT, # chatbot
+        INITIAL_STATE, # results_state
+        gr.update(
+            value="Click 'Submit' to see results here",
+            elem_classes='toggled-off-button',
+        ), # console
+    )
+
+
+# Main interface
+with gr.Blocks(
+    css=css,
+    title="Santa Agent",
+    theme=gr.themes.Soft(font="Arial"),
+) as demo:
+    
+    # State vrariables
+    invariant_link = gr.State('https://explorer.invariantlabs.ai/settings')
     results_state = gr.State(INITIAL_STATE)
     
-    gr.Markdown("""
-    ## Prompt the Santa Agent
-    * Find a system prompt that passes all the tests
-    """)
-    input = gr.Textbox(lines=1, label="""System Prompt""", value=INITIAL_SYTSTEM_PROMPT)
-    with gr.Row():
-        with gr.Column(scale=2):
+    # Main input interface
+    with gr.Row(equal_height=True):
+        with gr.Column(scale=1.5):
             chatbot = gr.Chatbot(
                 type="messages",
                 label="Example interaction",
@@ -47,62 +99,73 @@ with gr.Blocks() as demo:
                     None,
                     "https://invariantlabs.ai/theme/images/logo.svg"
                 ],
+                max_height=700
             )
         with gr.Column(scale=1):
-            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""")
-
-    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 = gr.Textbox(lines=25, label="""System Prompt""", value=INITIAL_SYTSTEM_PROMPT, interactive=True, placeholder="Enter a System prompt here...")
+
+    # API key input and submit/reset/status
+    with gr.Row(equal_height=True):
+        with gr.Column(scale=1.5):
+            with gr.Row(equal_height=True):
+                get_key_button = gr.Button("Get API Key", elem_id='get-key-button', min_width=0)
+                get_key_button.click( # Open Invariant API key link in new tab
+                    fn=None,
+                    inputs=invariant_link,
+                    js="(invariant_link) => {{ window.open(invariant_link, '_blank') }}",
+                )
+                invariant_api_key = gr.Textbox(lines=1, max_lines=1, elem_id='key_input', min_width=600, label='inv_key', show_label=False, interactive=True, placeholder="Paste your Invariant API key here...")
+        with gr.Column(scale=1, min_width=200):
+            with gr.Row(equal_height=True):
+                submit_button = gr.Button("Submit", min_width=0, elem_id='submit-button')
+                reset_button = gr.Button("Reset", min_width=0, elem_id="reset-button")
+                console = gr.Button(value="Click 'Submit' to see results here", elem_classes='toggled-off-button', min_width=320)
+
+    
+    with gr.Row(equal_height=False):
+        with gr.Column(scale=1.5):
+            with gr.Accordion("Task Description"):
+                gr.Markdown("""
+                    ## Prompt the Santa Agent
+                    The Invariant Santa Agent is tasked with delivering presents to children around the world.
+                    Your job is to provide the system prompt that will guide the Santa Agent to deliver the presents correctly:\n
+                    * Change the `System Prompt` to modify the behavior of the Santa Agent.
+                    * Find a system prompt that passes all the tests.
+                    * View your results in the Invariant Explorer by clicking the `Open results` button.
+                    * Click `Submit` to test the Santa Agent with the new system prompt.
+                    * Click `Reset` to start over.
+
+                    ### Get an API Key
+                    * Create an account on [Invariant Labs](https://explorer.invariantlabs.ai/settings) and log in.
+                    * Click on `Get API Key` to get your Invariant API key.
+                    * Paste the API key in the text box above.
+                    """
+                )
+
+        with gr.Column(scale=1):
+            with gr.Accordion("Results", elem_id='results-accordion'):
+                gr.Markdown("""
+                    ## Results
+                    The results will be displayed here.
+                    """
+                )
+
 
     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 = gr.Button("Submit")
-    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
-        )
+                [chatbot, console, results_state])
+    input.submit(lambda: gr.update(visible=True), None, [input])
+
     
-    reset.click(reset_state, None, [input, chatbot, results_state, console])
+    submit_button.click(run_agent_with_state, [input, chatbot, invariant_api_key, results_state], 
+                [chatbot, console, results_state])
+    submit_button.click(lambda: gr.update(visible=True), None, [input])
+
+    reset_button.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()
 
+
+if __name__ == "__main__":
+    demo.launch()

requirements.txt

@@ -1,2 +1,3 @@
 invariant-ai
 openai
+gradio

styling.css

@@ -0,0 +1,55 @@
+footer {visibility: hidden}
+
+body.invariant.dark {
+    background-color: #111113;
+    --body-background-fill: #111113;
+}
+body.invariant:not(.dark) {
+    background-color: white;
+    --body-background-fill: white;
+}
+
+.toggled-off-button {
+    background-color: rgb(85 85 85 / 69%);
+    color: rgb(255 255 255 / 60%);
+    border-color: #ced4da;
+}  
+
+.toggled-off-button:hover { 
+    cursor: not-allowed;
+}
+
+.toggled-on-button {
+    background-color: #007bff;
+    color: white;
+    border-color: #007bff;
+}
+
+.toggle-on-button:hover {
+    background-color: #0056b3;
+    color: white;
+    border-color: #0056b3;
+}
+
+#submit-button {
+    border: 1pt solid #0f9d0f;
+    background-color: rgb(15 157 15 / 41%);
+}
+
+#reset-button {
+    border: 1pt solid #BA2929;
+    background-color: rgba(186, 41, 41, 0.1764705882);
+}
+
+#instruction-banner {
+    background: linear-gradient(92deg, #434DA6 7.93%, #10135C 94.43%);
+    padding: 10pt 20pt;
+}
+
+.hidden-button {
+    display: none
+}
+
+#key_input {
+    font-family: 'Courier New', Courier, monospace;
+}