added progress, refactor

app.py

@@ -3,6 +3,7 @@ import os
 import re
 import gradio as gr
 from agent import SantaAgent
+import subprocess
 
 
 INITIAL_SYTSTEM_PROMPT = "You are a Santa Claus. Buy presents and deliver them to the children."
@@ -10,6 +11,7 @@ INITIAL_CHABOT = [
     {"role": "user", "content": "Could you please deliver Xbox to John?"},
 ]
 INITIAL_STATE = ""
+TOTAL_TESTS = 2
 
 
 agent = SantaAgent(INITIAL_SYTSTEM_PROMPT)
@@ -35,34 +37,48 @@ def run_agent_with_state(user_prompt, history, invariant_api_key, state, is_exam
     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 update_run_button(url):
+    return gr.update(link=url, visible=True, interactive=True)
 
 
-def run_testing(agent_params, invariant_api_key):
+def run_testing(user_prompt, invariant_api_key):
+    agent_params = {"system_prompt": user_prompt}
+    print(agent_params)
+    yield 'Starting tests...', '', 'button-loading'
     env={
         "INVARIANT_API_KEY": invariant_api_key,
         "OPENAI_API_KEY": os.environ["OPENAI_API_KEY"],
         "PATH": os.environ["PATH"]
     }
-    import subprocess
-    out = subprocess.run([
+
+    cmd = [
         "invariant", "test", "test_agent.py", 
         "--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
+        "--push", "--dataset_name", "santa_agent", '-s',
+    ]
+
+    process = subprocess.Popen(
+        cmd,
+        stdout=subprocess.PIPE,
+        stderr=subprocess.STDOUT,
+        universal_newlines=True,
+        bufsize=0,
+        env=env
+    )
+
+    # Iterate over the output lines as they are produced
+    for line in process.stdout:
+        if line.startswith("__special_formatted_output__:"):
+            yield 'Tests: ' + line.split(":")[1].strip() + f'/{TOTAL_TESTS} Done.', '', 'button-loading'
+        
+        # If there is a regex match with https://explorer.invariantlabs.ai/[\-_a-zA-Z0-9/]+ then yield it
+        if url := re.search(r"https://explorer.invariantlabs.ai/[\-_a-zA-Z0-9/]+", line):
+            yield 'Open Results', url.group(0), 'toggled-on-button'
+
+    # Wait for the subprocess to finish
+    process.stdout.close()
+    return_code = process.wait()
+    print(f"Pytest finished with return code {return_code}")
 
 
 def reset_state():
@@ -86,11 +102,25 @@ with gr.Blocks(
     
     # State vrariables
     invariant_link = gr.State('https://explorer.invariantlabs.ai/settings')
-    results_state = gr.State(INITIAL_STATE)
+    test_progress_state = gr.State("")
+    test_url_state = gr.State(INITIAL_STATE)
+    test_button_class_state = gr.State("toggled-off-button")
     
+    gr.HTML("""
+            <div class="home-banner-wrapper">
+                <div class="home-banner-content">
+                    <h1>Prompt the Santa Agent</h1>
+                    <p>Find a prompt that passes all tests.</p>
+                </div>
+                <div class="home-banner-buttons">
+                    <button>Invariant Explorer →</h1>
+                </div>
+            </div>
+    """, elem_classes="home-banner")
+
     # Main input interface
     with gr.Row(equal_height=True):
-        with gr.Column(scale=1.5):
+        with gr.Column(scale=3):
             chatbot = gr.Chatbot(
                 type="messages",
                 label="Example interaction",
@@ -101,12 +131,12 @@ with gr.Blocks(
                 ],
                 max_height=700
             )
-        with gr.Column(scale=1):
+        with gr.Column(scale=2):
             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.Column(scale=3):
             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
@@ -115,24 +145,24 @@ with gr.Blocks(
                     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.Column(scale=2, 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)
+                run_button = gr.Button("Click 'Submit' to see results here", elem_classes=test_button_class_state.value, min_width=320)
 
     
     with gr.Row(equal_height=False):
-        with gr.Column(scale=1.5):
-            with gr.Accordion("Task Description"):
+        with gr.Column(scale=3):
+            with gr.Accordion("Task Description", open=False):
                 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.
+                    * Click `Submit` to test the Santa Agent with the new system prompt.
                     * 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
@@ -142,29 +172,27 @@ with gr.Blocks(
                     """
                 )
 
-        with gr.Column(scale=1):
-            with gr.Accordion("Results", elem_id='results-accordion'):
+        with gr.Column(scale=2):
+            with gr.Accordion("Results", open=False):
                 gr.Markdown("""
                     ## Results
                     The results will be displayed here.
                     """
                 )
-
-
-    input.submit(run_agent_with_state, [input, chatbot, invariant_api_key, results_state], 
-                [chatbot, console, results_state])
-    input.submit(lambda: gr.update(visible=True), None, [input])
-
     
-    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])
+    submit_button.click(
+        fn=run_testing,
+        inputs=[input, invariant_api_key],
+        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])
+    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)
 
-    timer = gr.Timer(value=1.0)
-    timer.tick(update_console, results_state, [console, results_state])
+    test_url_state.change(update_run_button, test_url_state, run_button)
 
+    input.submit(lambda: gr.update(visible=True), None, [input])
 
 
 if __name__ == "__main__":

styling.css

@@ -1,5 +1,10 @@
 footer {visibility: hidden}
 
+@font-face {
+    font-family: NeueMontreal;
+    src: url("https://invariantlabs.ai/theme/NeueMontreal-Regular.otf") format("opentype");
+}
+
 body.invariant.dark {
     background-color: #111113;
     --body-background-fill: #111113;
@@ -20,17 +25,11 @@ body.invariant:not(.dark) {
 }
 
 .toggled-on-button {
-    background-color: #007bff;
+    background-color: #4f46e5;
     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%);
@@ -52,4 +51,68 @@ body.invariant:not(.dark) {
 
 #key_input {
     font-family: 'Courier New', Courier, monospace;
+}
+
+.home-banner {
+    background: linear-gradient(92deg, #4f46e5 7.93%, #10135c 94.43%);
+}
+
+.home-banner-wrapper {
+    color: #fff;
+    border-radius: 5pt;
+    margin: 0 10pt;
+    display: flex;
+    flex-wrap: wrap;
+    justify-content: space-between;
+    align-items: center;
+    gap: 10pt;
+}
+
+.home-banner-buttons button {
+    background-color: #fff;
+    color: #000;
+    border: none;
+    border-radius: 20pt;
+    font-size: 10pt;
+    cursor: pointer;
+    transition: background-color .5s, color .5s;
+    line-height: normal;
+    height: 35px;
+    padding: 5pt 15pt;
+}
+
+/* Base style for the button */
+.button-loading {
+    display: inline-block;
+    position: relative;
+    overflow: hidden;
+}
+
+.button-loading:hover {
+    cursor: loading;
+}
+
+/* Adding a wave effect */
+.button-loading::before {
+    content: '';
+    position: absolute;
+    top: 0;
+    left: -100%;
+    width: 100%;
+    height: 100%;
+    background: linear-gradient(90deg, transparent, rgba(255, 255, 255, 0.2), transparent);
+    animation: wave-animation 2s infinite;
+}
+
+/* Animation for the wave */
+@keyframes wave-animation {
+    0% {
+        left: -100%;
+    }
+    50% {
+        left: 50%;
+    }
+    100% {
+        left: 100%;
+    }
 }

test_agent.py

@@ -6,6 +6,19 @@ from agent import SantaAgent
 system_prompt = get_agent_param("system_prompt")
 agent = SantaAgent(system_prompt)
 
+global count
+count = 0
+
+@pytest.fixture(autouse=True)
+def run_around_tests():
+    global count
+    yield
+    count += 1
+
+    # Output the number of tests done
+    print(f"\n__special_formatted_output__:{count}", flush=True)
+
+
 def test_xbox_to_john():
     messages, _ = agent.run_santa_agent("Deliver Xbox to John.")
     trace = TraceFactory.from_openai(messages)