app.py

import os
if not os.getenv("OPENAI_API_KEY"):
    raise ValueError("OPENAI_API_KEY must be set")

if os.getenv("WITH_REACHY"):
    WITH_REACHY = True
else:
    WITH_REACHY = False

if os.getenv("IN_DOCKER"):
    IN_DOCKER = True
else:
    IN_DOCKER = False


import gradio as gr
import uuid
from pictionagent import AgentStatus, ItemTopic, PictionagentManager
import logging
logging.basicConfig(level=logging.INFO)

ELEMENTS_UPDATE_PERIOD = 0.5 # seconds

pictionagent_manager = PictionagentManager(fake_robot=not WITH_REACHY)

#Gradio update functions
def connect_user():
    user_id = str(uuid.uuid4())
    (nb_connected_users, connected_users, nb_connected_users_game, user_id_game, user_score_game, leaderboard) = update_users(user_id)
    return user_id, user_id, nb_connected_users, connected_users, nb_connected_users_game,  user_id_game, user_score_game, leaderboard

def tick_check_update(history_state: int, hidden_for_history_changes: gr.Textbox,
                      hidden_for_guess_interactivity_changes: gr.Textbox,
                      hidden_for_game_interactivity_changes: gr.Textbox):
    # Check history state change
    if history_state != len(pictionagent_manager.agent_history):
        history_state = len(pictionagent_manager.agent_history)
        hidden_for_history_changes = str(uuid.uuid4())

    # Game elements interactivity
    if not pictionagent_manager.async_init_done:
        hidden_for_game_interactivity_changes = "False"
    elif pictionagent_manager.game_running:
        hidden_for_game_interactivity_changes = "False"
    else:
        hidden_for_game_interactivity_changes = "True"

    # Guess elements interactivity
    if not pictionagent_manager.async_init_done:
        hidden_for_guess_interactivity_changes = "False"
    elif (pictionagent_manager.agent_status == AgentStatus.DRAWING or pictionagent_manager.agent_status == AgentStatus.WAITING_FOR_GUESS) and not pictionagent_manager.is_drawing_guessed:
        hidden_for_guess_interactivity_changes = "True"
    else:
        hidden_for_guess_interactivity_changes = "False"


    return (pictionagent_manager.drawing_image,
            pictionagent_manager.generated_source_image,
            pictionagent_manager.displayed_item_to_guess,
            pictionagent_manager.agent_status.value,
            pictionagent_manager.nb_drawing_points,
            pictionagent_manager.drawing_duration,
            history_state,
            hidden_for_history_changes,
            hidden_for_game_interactivity_changes,
            hidden_for_guess_interactivity_changes)


def update_history():
    return pictionagent_manager.agent_history

def change_topic(topic:str):
    pictionagent_manager.current_topic = ItemTopic(topic)

def set_nb_points(nb_points: int):
    pictionagent_manager.nb_drawing_points = nb_points

def set_draw_duration(draw_duration: float):
    pictionagent_manager.drawing_duration = draw_duration   

def update_game_elements_interactivity(hidden_for_game_interactivity_changes: gr.Textbox):
    if hidden_for_game_interactivity_changes == "True":
        return gr.update(interactive=True), gr.update(interactive=True), gr.update(interactive=True), gr.update(interactive=True), gr.update(interactive=True)
    else:
        return gr.update(interactive=False), gr.update(interactive=False), gr.update(interactive=False), gr.update(interactive=False), gr.update(interactive=False)

def update_guess_elements_interactivity(hidden_for_guess_interactivity_changes: gr.Textbox):
    if hidden_for_guess_interactivity_changes == "True":
        return gr.update(interactive=True), gr.update(interactive=True)
    else:
        return gr.update(value="", interactive=False), gr.update(interactive=False)

async def guess_drawing(guess_input: str, user_id: str):
    if not pictionagent_manager.async_init_done:
        return
    elif (pictionagent_manager.agent_status != AgentStatus.DRAWING and pictionagent_manager.agent_status != AgentStatus.WAITING_FOR_GUESS) or pictionagent_manager.is_drawing_guessed:
        return

    await pictionagent_manager.try_guess_drawing(guess_input, user_id)

def reset_game():
    pictionagent_manager.reset_game()
    return [], pictionagent_manager.agent_status, pictionagent_manager.drawing_image, pictionagent_manager.generated_source_image, pictionagent_manager.displayed_item_to_guess

def update_users(user_id: str):
    leaderboard = ""
    for i in range(0, min(len(pictionagent_manager.leaderboard), 10)):
        leaderboard += str(i + 1) + ") " + pictionagent_manager.get_user_display_name(pictionagent_manager.leaderboard[i].id) + " : " + str(pictionagent_manager.leaderboard[i].score) + "\n"

    if user_id is None:
        return 0, "", 0, "", 0, leaderboard
    
    pictionagent_manager.on_user_ping(user_id)
    if pictionagent_manager.connected_users:
        output = ""
        for id in pictionagent_manager.connected_users.keys():
            output += pictionagent_manager.get_user_display_name(id) + "\n"
        return len(pictionagent_manager.connected_users), output, len(pictionagent_manager.connected_users), pictionagent_manager.get_user_display_name(user_id), pictionagent_manager.connected_users[user_id].score, leaderboard
    else:
        return 0, "", 0, "", 0, leaderboard

def update_name(name_box: str, user_id: str):
    pictionagent_manager.set_user_name(user_id, name_box)
    return update_users(user_id)
#endregion

if IN_DOCKER:
    vnc_url = "/vnc/vnc.html?autoconnect=1"
else:
    vnc_url = "http://localhost:6080/vnc.html?autoconnect=1"

reachy_iframe_html = f"""
<div style="overflow:hidden; width:350px; height:500px; pointer-events: none; user-select: none;">
  <iframe src="{vnc_url}"
          tabindex="-1"
          style="
            transform: scale(1) translate(-45%, -15%);
            transform-origin: top left;
            width: 1024px;
            height: 768px;
            border: none;">
  </iframe>
</div>
"""
# Gradio UI
with gr.Blocks() as demo:
    if WITH_REACHY:
        gr.Markdown("# Pictionagent x Reachy")
    else:
        gr.Markdown("# Pictionagent")
    user_state = gr.State()

    with gr.Tab("Game"):
        with gr.Row():
            with gr.Column(scale=1):
                nb_connected_users_game = gr.Textbox(label="Connected users", value=0, interactive=False)
            with gr.Column(scale=1):
                user_id_game = gr.Textbox(label="You", value="", interactive=False)
            with gr.Column(scale=1):
                user_score_game = gr.Textbox(label="Your score", value="", interactive=False)
        with gr.Row():
            with gr.Column(scale=1):
                topic_dropdown = gr.Dropdown(choices=[e.value for e in ItemTopic],
                                            label="Pick a topic")
            with gr.Column(scale=1):
                start_btn = gr.Button("Start round")
                clear_btn = gr.Button("Reset")
        with gr.Row():
            with gr.Column(scale=2 if WITH_REACHY else 1):
                agent_status_display = gr.Textbox(label="Agent Status", value=pictionagent_manager.agent_status.value)
                chatbot = gr.Chatbot(pictionagent_manager.agent_history, type="messages", show_label=False)
            if WITH_REACHY:
                with gr.Column(scale=2):
                    gr.HTML(reachy_iframe_html)
            with gr.Column(scale=3 if WITH_REACHY else 1):
                item_to_guess = gr.Textbox(label="Item to guess", value="", interactive=False)
                with gr.Row():
                    drawing_image_display = gr.Image(label="Drawn Image", type="pil", interactive=False, show_label=False)
                    source_image_display = gr.Image(label="Source Image", type="pil", interactive=False, show_label=False)
                gr.Markdown("### Your Guess")
                with gr.Row():
                    guess_input = gr.Textbox(label="Guess", value="", show_label=False)
                    guess_btn = gr.Button("Submit")
        
        hidden_for_history_changes = gr.Textbox(value="", visible=False, interactive=False)
        history_state = gr.State(0)
        hidden_for_guess_interactivity_changes = gr.Textbox(value="", visible=False, interactive=False)
        hidden_for_game_interactivity_changes = gr.Textbox(value="", visible=False, interactive=False)

    with gr.Tab("Connection Infos"):
        with gr.Column(scale=1):
            id_box = gr.Textbox(label="Your ID", interactive=False)
            name_box = gr.Textbox(label="Your name")
            name_btn = gr.Button("Set/Update name")
        with gr.Column(scale=1):
            nb_connected_users = gr.Textbox(label="Number of connected users", value=0, interactive=False)
            connected_users = gr.Textbox(label="Connected users", value="", lines=5, max_lines=5, interactive=False)

    with gr.Tab("Game config"):
        nb_points_slider = gr.Slider(label="Drawing points", value=100, minimum=20, maximum=100, step=10)
        draw_duration_slider = gr.Slider(label="Duration between points", value=0.05, minimum=0.02, maximum=0.2, step=0.01)

    with gr.Tab("Leaderboard"):
        leaderboard = gr.Textbox(label="Leaderboard", value="", lines=10, max_lines=10, interactive=False)

    # Game tab
    topic_dropdown.change(change_topic, inputs=[topic_dropdown], outputs=[], api_name=False)
    start_btn.click(pictionagent_manager.start_game, inputs=[], outputs=[], api_name=False)
    guess_btn.click(guess_drawing, inputs=[guess_input, user_state], outputs=[], api_name=False)
    guess_input.submit(guess_drawing, inputs=[guess_input, user_state], outputs=[], api_name=False)

    clear_btn.click(reset_game, inputs=[], outputs=[chatbot, agent_status_display, drawing_image_display, source_image_display, item_to_guess], api_name=False)

    # Timer triggers the update function with given period
    # + workaround (with hidden textboxes) to update elements only on value change (not on each time tick)
    timer = gr.Timer(value=ELEMENTS_UPDATE_PERIOD)
    timer.tick(tick_check_update,
               inputs=[history_state, hidden_for_history_changes, hidden_for_game_interactivity_changes, hidden_for_guess_interactivity_changes],
               outputs=[drawing_image_display, source_image_display, item_to_guess, agent_status_display, nb_points_slider, draw_duration_slider, history_state, hidden_for_history_changes, hidden_for_game_interactivity_changes, hidden_for_guess_interactivity_changes],
               api_name=False)
    hidden_for_history_changes.change(update_history, outputs=chatbot, api_name=False)
    hidden_for_game_interactivity_changes.change(update_game_elements_interactivity,
                                                 inputs=[hidden_for_game_interactivity_changes], outputs=[topic_dropdown, start_btn, clear_btn, nb_points_slider, draw_duration_slider], api_name=False)
    hidden_for_guess_interactivity_changes.change(update_guess_elements_interactivity,
                                                  inputs=[hidden_for_guess_interactivity_changes], outputs=[guess_input, guess_btn],
                                                  api_name=False)

    # Connection tab
    connection_timer = gr.Timer(value=1)
    connection_timer.tick(update_users, inputs=[user_state], outputs=[nb_connected_users, connected_users, nb_connected_users_game, user_id_game, user_score_game, leaderboard], api_name=False)
    name_btn.click(update_name, inputs=[name_box, user_state], outputs=[nb_connected_users, connected_users, nb_connected_users_game, user_id_game, user_score_game, leaderboard], api_name=False)

    # Game config tab
    nb_points_slider.change(set_nb_points, inputs=[nb_points_slider], outputs=[], api_name=False)
    draw_duration_slider.change(set_draw_duration, inputs=[draw_duration_slider], outputs=[], api_name=False)

    # Register the init_agent to run before interface is live
    demo.load(fn=pictionagent_manager.async_init, inputs=[], outputs=[], queue=False, api_name=False)
    demo.load(fn=connect_user, inputs=[], outputs=[user_state, id_box, nb_connected_users, connected_users, nb_connected_users_game, user_id_game, user_score_game, leaderboard], api_name=False)

# When using the docker image, we use a nginx reverse proxy to access the gradio app
# So we need to set the root_path to "/gradio" to route it properly via nginx
if IN_DOCKER:
    root_path = "/gradio"
else:
    root_path = None
demo.launch(server_name="0.0.0.0", server_port=7860, root_path=root_path)