import gradio as gr
import requests
from tqdm import tqdm
from MonsterAPIClient import MClient
from MonsterAPIClient import MODELS_TO_DATAMODEL
client = MClient()



# Available models list
EXCLUSION_LIST = ['mpt-30B-instruct']
available_models = list(set(list(MODELS_TO_DATAMODEL.keys())) - set(EXCLUSION_LIST))

def generate_model_output(model, input_text):
    try:
        response = client.get_response(model, {"prompt": input_text})
        output = client.wait_and_get_result(response['process_id'])
        return output
    except Exception as e:
        return f"Error occurred: {str(e)}"

# Gradio interface function
def generate_output(selected_models, input_text, available_models=available_models):
    outputs = {}
    for model in tqdm(selected_models):
        outputs[model] = generate_model_output(model, input_text)
    ret_outputs = []
    for model in available_models:
        if model not in outputs:
            ret_outputs.append("Model not selected!")
        else:
            ret_outputs.append(outputs[model].replace("\n", "<br>"))

    return ret_outputs
    
output_components = [gr.outputs.Textbox(label=model) for model in available_models]

checkboxes = gr.inputs.CheckboxGroup(available_models , label="Select models to generate outputs:")
textbox = gr.inputs.Textbox()

# Gradio Interface
input_text = gr.Interface(
    fn=generate_output,
    inputs=[checkboxes, textbox],
    outputs=output_components,
    live=False,
    capture_session=True,
    title="Monster API LLM Output Comparison.",
    description="Generate outputs from selected models using Monster API.",
    css="body {background-color: black}"
)

# Launch the Gradio app
input_text.launch()