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): 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"] } 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("""

Prompt the Santa Agent

Find a prompt that passes all tests.

""", 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()