import gradio as gr
import requests
import json

def call_api(prompt, history_json, system_prompt, temperature, max_new_tokens, top_p, repetition_penalty):
    url = "https://seawolf2357-nvidia-llama-fastapi.hf.space/generate/"
    try:
        # JSON 형태의 history를 파싱 시도
        history = json.loads(history_json)
    except json.JSONDecodeError:
        return "History input must be valid JSON. E.g., [[\"User\", \"Bot response\"], [\"User2\", \"Response2\"]]"

    data = {
        "prompt": prompt,
        "history": history,
        "system_prompt": system_prompt,
        "temperature": temperature,
        "max_new_tokens": max_new_tokens,
        "top_p": top_p,
        "repetition_penalty": repetition_penalty
    }

    try:
        response = requests.post(url, json=data)
        response.raise_for_status()  # Raises HTTPError for bad requests
        return response.json()['response']
    except requests.exceptions.RequestException as e:
        return f"An error occurred: {str(e)}"

def main_interface():
    with gr.Blocks() as demo:
        with gr.Row():
            prompt = gr.Textbox(label="Prompt", placeholder="Enter your prompt here...")
            history = gr.TextArea(label="History", placeholder="Enter history as JSON array...", lines=3)
            system_prompt = gr.Textbox(label="System Prompt", placeholder="반드시 한글로 답변하라")
            temperature = gr.Number(label="Temperature", value=0.3)
            max_new_tokens = gr.Number(label="Max New Tokens", value=1048)
            top_p = gr.Number(label="Top P", value=0.95)
            repetition_penalty = gr.Number(label="Repetition Penalty", value=1.0)
            submit_button = gr.Button("Generate")

        submit_button.click(
            fn=call_api,
            inputs=[prompt, history, system_prompt, temperature, max_new_tokens, top_p, repetition_penalty],
            outputs=[gr.Textbox(label="Response")]
        )

    return demo

if __name__ == "__main__":
    demo = main_interface()
    demo.launch()