app.py

import json
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."
INITIAL_CHABOT = [
    {"role": "user", "content": "Could you please deliver Xbox to John?"},
]
INITIAL_STATE = ""
TOTAL_TESTS = 2


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_run_button(url):
    return gr.update(link=url, visible=True, interactive=True)


def run_testing(user_prompt, invariant_api_key):
    if not invariant_api_key.startswith("inv"):
        return "Please enter a valid Invariant API key to get the score!", '', 'toggled-off-button'

    agent_params = {"system_prompt": user_prompt}

    yield 'Starting tests...', '', 'button-loading'
    env={
        "INVARIANT_API_KEY": invariant_api_key,
        "OPENAI_API_KEY": os.environ["OPENAI_API_KEY"],
        "PATH": os.environ["PATH"]
    }

    cmd = [
        "invariant", "test", "test_agent.py", 
        "--agent-params", json.dumps(agent_params),
        "--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():
    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')
    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=3):
            chatbot = gr.Chatbot(
                type="messages",
                label="Example interaction",
                value=INITIAL_CHABOT,
                avatar_images=[
                    None,
                    "https://invariantlabs.ai/theme/images/logo.svg"
                ],
                max_height=700
            )
        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=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
                    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=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")
                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=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 `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=2):
            with gr.Accordion("Results", open=False):
                gr.Markdown("""
                    ## Results
                    The results will be displayed here.
                    """
                )
    
    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)

    test_url_state.change(update_run_button, test_url_state, run_button)

    input.submit(lambda: gr.update(visible=True), None, [input])


if __name__ == "__main__":
    demo.launch()