import gradio as gr
from huggingface_hub import InferenceClient

# 1) Initialize your HF client
client = InferenceClient("HuggingFaceH4/zephyr-7b-beta")

# 2) A pure, one-shot function
def respond(
    system_message: str,
    user_prompt: str,
    max_tokens: int,    # ← add this parameter
    temperature: float,
    top_p: float
) -> str:


    messages = [
        {"role": "system", "content": system_message},
        {"role": "user",   "content": user_prompt}
    ]

    resp = client.chat_completion(
        messages,
        max_tokens=max_tokens,
        temperature=temperature,
        top_p=top_p,
        stream=False    # <-- non-streaming
    )

    return resp.choices[0].message.content.strip()

# 3) Create a non-streaming Interface
demo = gr.Interface(
    fn=respond,
    inputs=[
        gr.Textbox("You are a helpful and friendly assistant for children. Your capabilities include answering questions about children's books, characters, and events, and generating engaging, imaginative stories. It is essential that all your responses and generated stories are safe, appropriate, and positive for children. If a question or request is inappropriate, or if it is not about children's books or generating a safe story, please politely explain that you can only help with safe requests related to children's books or stories.", label="System message"),
        gr.Textbox(placeholder="Type your question…", label="User prompt"),
        gr.Slider(1, 2048, 512, label="Max tokens"),
        gr.Slider(0.1, 4.0, 0.7, label="Temperature"),
        gr.Slider(0.1, 1.0, 0.95, label="Top-p"),
    ],
    outputs=gr.Textbox(label="Assistant reply"),
    title="StoryNest Chat (One-Shot)",
    flagging_mode="never",
)

# 4) Enable the REST API and then launch
demo.queue(api_open=True)
if __name__ == "__main__":
    demo.launch()