import gradio as gr
import outlines
import os, json
from outlines.models import openai
from outlines.generate import choice


model_list = {
    "Gemma-2B": "google/gemma-2b-it",
    "Qwen2.5-7B": "Qwen/Qwen2.5-7B-Instruct-Turbo"
}

def create_prompt_message(message: str) -> str:
    return f"""
        You are an agent that classifies restaurant-related user messages.

        Your job is to classify the following message into one of two categories: "Reservation" or "Others".

        You MUST respond EXACTLY with one of these two JSON objects:

        {{"result": "Reservation"}} or {{"result": "Others"}}

        Do not add any explanation or other content.

        Message: {message}
        """.strip()

def inference(model_dropdown, message):
    # 根據 model_name 決定要用哪個模型
    model_id = model_list[model_dropdown]
    model = openai(
        model_id,
        api_key=os.getenv("TOGETHER_API_KEY"),
        base_url="https://api.together.xyz/v1"
    )
    classifier = choice(model, ["Reservation", "Others"])

    # 下面是原本分類流程
    prompt = create_prompt_message(message)
    try:
        result = classifier(prompt)
        if result == 'Reservation':
            return "Reservation request！"
        elif result == 'Others':
            return "Non-reservation request！"
        else:
            return "不明意圖。"
    except Exception as e:
        print(f"[Error!]: {e}")
        return "Error!"


with gr.Blocks() as demo:
    gr.Markdown("## 🍴🥂 LLM訂位意圖分類器")
    gr.Markdown("判斷輸入訊息是否與訂位有關")

    with gr.Row():
        model_dropdown = gr.Dropdown(choices=list(model_list.keys()), label="選擇模型", value="Gemma-2B")
    with gr.Row():
        input_box=gr.Textbox(placeholder="請輸入顧客訊息...", label="輸入訊息")
        output_box=gr.Textbox(label="分類結果")
    with gr.Row():
        inference_btn = gr.Button("開始分類")

    inference_btn.click(
        fn=inference,
        inputs=[model_dropdown, input_box],
        outputs=output_box
    )

demo.launch()