app.py

import gradio as gr
from chatbot_simulator import ChatbotSimulation
from task_specific_data_population import DataPopulation
from datasets import load_dataset
import json_repair
import random

openai_api_key = os.getenv("OPENAI_API_KEY")


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

    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'])

        synthetic_tasks = load_dataset(sitemap_url, "tasks", split='train')
        random_index = random.randint(0, len(synthetic_tasks) - 1)  # Generate a random index
        random_row = synthetic_tasks[random_index]
        task = random_row['tasks']
        solution = random_row['steps']
        user_data = random_row['attributes']['user_data']

        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)
    return response


# 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')

    # 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
    )

# Launch the app
demo.launch()