import gradio as gr
from huggingface_hub import InferenceClient

"""
For more information on `huggingface_hub` Inference API support, please check the docs: https://huggingface.co/docs/huggingface_hub/v0.22.2/en/guides/inference
"""
# client = InferenceClient("ibm-granite/granite-3.1-2b-instruct") # specifically for long contexts.
client = InferenceClient("meta-llama/Llama-3.2-3B-Instruct") # too small context window
# client = InferenceClient("HuggingFaceH4/zephyr-7b-beta")

_sys_msg = """
You are an AI language model specialized in mindful communication in English. I prompt you with text, that you should edit and re-write. Your responsibilities are to:


1. **Politeness & Formality:** Rewrite the provided input text to be polite and formal.
2. **Disambiguation:** Clarify any ambiguous terms or phrases to ensure the message is clear and unambiguous.
3. **Explanation:** After revising the text, provide a detailed explanation of the changes you made and the reasons behind each change.


Ensure that all text you re-write is clear and respectful and maintain the original intent of the message while enhancing its clarity and professionalism."""

_ass_msg_start = """
Welcome to the interview. I want to write a short biography about you and need some input from your side. 
Can you please start by stating your name and talking about your early childhood?
"""

def respond(
    message,
    history: list[tuple[str, str]],
):
    messages = [{"role": "system", "content": _sys_msg},]
    
    for val in history:
        if val[0]:
            messages.append({"role": "user", "content": val[0]})
        if val[1]:
            messages.append({"role": "assistant", "content": val[1]})


    messages.append({"role": "user", "content": message})

    response = ""

    for message in client.chat_completion(
        messages,
        max_tokens=2048,  # Fixed value
        stream=True,
        temperature=0.5,  # Fixed value
        top_p=0.95,      # Fixed value
    ):
        token = message.choices[0].delta.content

        response += token
        yield response


demo  = gr.ChatInterface(
    respond,
    title = "Mindful Communication Manager",
    description = """
    An AI system prompt helps re-formulate communication to be more polite and mindful. 
    """,
    examples=["Help me write an email with the following content: I think we are stuck and we should meet to sort out the next steps. Tue and Wed evening are good for me. in person better than online. let me know asap."]
    )


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