app.py

import gradio as gr
from chatbot_simulator import ChatbotSimulation
from huggingface_hub import HfApi, create_repo
from datasets import load_dataset
import json_repair
import random
import os
import re
os.environ["TOKENIZERS_PARALLELISM"] = "false"

openai_api_key = os.getenv("OPENAI_API_KEY")

task_completed = 0
task_completed_steps = None


def is_task_incomplete(example):
    return example["task_completed"] is None and example["task_completed_steps"] is None


def is_task_given_up(example):
    return example["task_completed"] == 0 and example["task_completed_steps"] == 0


def is_task_incomplete_idx(example, idx):
    return example["task_completed"] is None and example["task_completed_steps"] is None


def is_task_given_up_idx(example, idx):
    return example["task_completed"] == 0 and example["task_completed_steps"] == 0


class AppSimulator:
    def __init__(self, openai_api_key):
        self.simulation = None
        self.openai_api_key = openai_api_key
        self.app_name = None

    def initialize_simulator(self, sitemap_url, progress=gr.Progress(track_tqdm=True)):
        """Initialize the simulator with retries and elapsed time tracking."""
        synthetic_sitemap = load_dataset(sitemap_url, "sitemap", split='train')
        app_name, sitemap, page_details, user_state, system_data = None, None, None, None, None
        for row in synthetic_sitemap:
            if row['name'] == 'app_name':
                app_name = row['value']  # Use `eval` to convert the string to a list
            elif row['name'] == 'sitemap':
                sitemap = json_repair.loads(row['value'])
            elif row['name'] == 'page_details':
                page_details = json_repair.loads(row['value'])
            elif row['name'] == 'user_state':
                user_state = json_repair.loads(row['value'])
            elif row['name'] == 'system_data':
                system_data = json_repair.loads(row['value'])

        self.app_name = app_name
        human_result_url = "jjz5463/simulator_human_result"
        synthetic_tasks = load_dataset(human_result_url, app_name, split='train')

        incomplete_tasks = synthetic_tasks.filter(is_task_incomplete)
        give_up_tasks = synthetic_tasks.filter(is_task_given_up)
        if len(incomplete_tasks) > 0:
            incomplete_task = incomplete_tasks[0]
            task = incomplete_task["tasks"]
            solution = incomplete_task["steps"]
            user_data = incomplete_task["attributes"]["user_data"]
        elif len(give_up_tasks) > 0:
            give_up_task = give_up_tasks[0]
            task = give_up_task["tasks"]
            solution = give_up_task["steps"]
            user_data = give_up_task["attributes"]["user_data"]
        else:
            return "All tasks in this app have been completed!"

        self.simulation = ChatbotSimulation(
            app_name=app_name,
            site_map=sitemap,
            page_details=page_details,
            user_state=user_state,
            system_data=system_data,
            user_data=user_data,
            task=task,
            solution=solution,
            log_location=f'conversation_log_{app_name}.txt',
            openai_api_key=openai_api_key,
            agent='llm'
        )

        initial_message = self.simulation.start_conversation()
        progress.update("Initialization Successful")
        return initial_message  # Return the initial assistant message for chat

    def chat_interaction(self, user_input, history):
        """Handle one round of conversation."""
        return self.simulation.one_conversation_round(user_input)


# Initialize the simulator
simulator_app = AppSimulator(openai_api_key=openai_api_key)


def chat(user_input, history):
    """Chat handler that validates input and interacts with the simulator."""
    response = simulator_app.chat_interaction(user_input, history)

    # Initialize variables for task completion and steps

    # Define the pattern for matching the response
    pattern = r"Task completed! You took (\d+) steps\."
    match = re.match(pattern, response)

    global task_completed
    global task_completed_steps
    task_completed = 0
    task_completed_steps = None
    if match:
        task_completed = 1
        task_completed_steps = int(match.group(1))
        human_result_url = "jjz5463/simulator_human_result"
        app_name = simulator_app.app_name
        synthetic_tasks = load_dataset(human_result_url, app_name, split='train')
        incomplete_tasks = synthetic_tasks.filter(is_task_incomplete_idx, with_indices=True,)
        if len(incomplete_tasks) > 0:
            incomplete_task_index = incomplete_tasks['"__index"'][0]
            synthetic_tasks = synthetic_tasks.map(
                lambda example, idx: {
                    "task_completed": task_completed if idx == incomplete_task_index else example["task_completed"],
                    "task_completed_steps": task_completed_steps if idx == incomplete_task_index else example["task_completed_steps"],
                },
                with_indices=True,
            )
            # Push the updated dataset back to Hugging Face
            synthetic_tasks.push_to_hub(human_result_url)

    return response


def give_up():
    """Handle the Give-Up action by marking the first incomplete task as abandoned."""
    global task_completed, task_completed_steps
    task_completed = 0
    task_completed_steps = 0

    # Access the app_name from the simulator instance
    app_name = simulator_app.app_name
    if not app_name:
        return "Simulator has not been initialized with an app!"

    # Load the human result dataset
    human_result_url = "jjz5463/simulator_human_result"
    synthetic_tasks = load_dataset(human_result_url, app_name, split='train')

    # Find the first incomplete task
    incomplete_tasks = synthetic_tasks.filter(is_task_incomplete_idx, with_indices=True,)
    if len(incomplete_tasks) > 0:
        incomplete_task_index = incomplete_tasks['"__index"'][0]

        # Update the dataset to mark the task as abandoned
        synthetic_tasks = synthetic_tasks.map(
            lambda example, idx: {
                "task_completed": task_completed if idx == incomplete_task_index else example["task_completed"],
                "steps": task_completed_steps if idx == incomplete_task_index else example["steps"],
            },
            with_indices=True,
        )

        # Push the updated dataset back to Hugging Face
        synthetic_tasks.push_to_hub(human_result_url)
        return "Task marked as abandoned (Give-Up action)."
    else:
        return "No incomplete tasks found to abandon!"


# Gradio Interface using ChatInterface
with gr.Blocks(fill_height=True) as demo:
    gr.Markdown("## Simulator Setup")

    # Input fields for initialization
    sitemap_input = gr.Textbox(label="Sitemap", placeholder="Enter the Hugging Face link to sitemap... (eg.jjz5463/DoorDash_synthetic_sitemap)")
    initialize_button = gr.Button("Initialize Simulator")

    # Status block to display initialization progress with elapsed time
    status = gr.Textbox(label="Status", interactive=False)

    # Chat interface to handle user interactions
    chat_interface = gr.ChatInterface(fn=chat, type='messages')

    give_up_button = gr.Button("Give Up")

    # Define the callback function to initialize the simulator and update status
    def initialize_and_start_chat(sitemap):
        return simulator_app.initialize_simulator(sitemap)  # Use progress tracking

    # Set up the button click to initialize simulator and update status only
    initialize_button.click(
        fn=initialize_and_start_chat,
        inputs=[sitemap_input],
        outputs=status  # Update only the status block
    )

    # Set up the Give-Up button click to update the dataset
    give_up_button.click(
        fn=give_up,
        inputs=[],
        outputs=status  # Update the status with the give-up message
    )

# Launch the app
demo.launch()