import gradio as gr
from huggingface_hub import InferenceClient

# Initialize the InferenceClient with the model ID from Hugging Face
client = InferenceClient(model="HuggingFaceH4/zephyr-7b-beta")

# Latest updates (you can replace this with actual update information)
latest_updates = """
**Zephyr 7B Beta Chatbot - Latest Updates:**

* **Improved Error Handling:** The chatbot now provides clearer error messages if something goes wrong.
* **Enhanced System Message Input:** You can now provide multi-line system messages to guide the AI's behavior.
* **Optimized Temperature Range:** The temperature slider's range has been adjusted for better control over randomness.
* **Robust Chunk Handling:** The chatbot now handles streamed responses more reliably, even if some chunks are missing content. 
"""

def respond(
    message: str,
    history: list[tuple[str, str]],
    system_message: str,
    max_tokens: int,
    temperature: float,
    top_p: float,
):
    """
    Generates a response from the AI model based on the user's message and chat history.

    Args:
        message (str): The user's input message.
        history (list): A list of tuples representing the conversation history (user, assistant).
        system_message (str): A system-level message guiding the AI's behavior.
        max_tokens (int): The maximum number of tokens for the output.
        temperature (float): Sampling temperature for controlling the randomness.
        top_p (float): Top-p (nucleus sampling) for controlling diversity.

    Yields:
        str: The AI's response as it is generated.
    """

    # Prepare the conversation history for the API call
    messages = [{"role": "system", "content": system_message}]

    for user_input, assistant_response in history:
        messages.append({"role": "user", "content": user_input})
        messages.append({"role": "assistant", "content": assistant_response})

    # Add the latest user message to the conversation
    messages.append({"role": "user", "content": message})

    # Initialize an empty response
    streamed_response = ""

    try:
        # Generate a response from the model with streaming
        for response in client.chat_completion(
            messages=messages,
            max_tokens=max_tokens,
            stream=True,
            temperature=temperature,
            top_p=top_p,
        ):
            chunk = response.choices[0].delta.get("content", "")
            streamed_response += chunk
            yield streamed_response

    except Exception as e:
        yield f"**Error:** {str(e)}"


# Define the Gradio interface with the Blocks context
with gr.Blocks() as demo:
    chat_interface = gr.ChatInterface(
        fn=respond,
        additional_inputs=[
            gr.Textbox(
                value="You are a friendly and helpful assistant.",
                label="System message",
                lines=2
            ),
            gr.Slider(minimum=1, maximum=2048, value=512, step=1, label="Max new tokens"),
            gr.Slider(minimum=0.1, maximum=2.0, value=0.7, step=0.1, label="Temperature"),
            gr.Slider(
                minimum=0.1,
                maximum=1.0,
                value=0.95,
                step=0.05,
                label="Top-p (nucleus sampling)",
            ),
        ],
        title="Zephyr 7B Beta Chatbot",
        description="A customizable chatbot interface using Hugging Face's Zephyr 7B Beta model and Inference API.",
    )

    # Add the "Show Updates" button and output area
    with gr.Row():
        updates_button = gr.Button("Show Latest Updates")
        updates_output = gr.Markdown(visible=False)

    # Define the button's click event (now inside the Blocks context)
    updates_button.click(
        fn=lambda: latest_updates,
        outputs=updates_output,
        show_progress=False
    )


# Launch the Gradio interface 
if __name__ == "__main__":
    demo.launch()