import gradio as gr
from huggingface_hub import InferenceClient
from huggingface_hub import login
import os
import time 

L_ogin = os.getenv('token')



"""
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
"""




import time

def respond(
    message,
    history: list[tuple[str, str]],
    model_id,
    system_message,
    max_tokens,
    temperature,
    top_p,
):
    #client = InferenceClient(token=L_ogin,timeout = 30, model=kheops_models.get(model_id))
    client = InferenceClient()
    messages = [{"role": "system", "content": system_message}]

    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=max_tokens,
        stream=True,
        temperature=temperature,
        top_p=top_p,
    ):
        token = message.choices[0].delta.content
        if token is None:
            token = ""
        response += token
        yield response


"""
For information on how to customize the ChatInterface, peruse the gradio docs: https://www.gradio.app/docs/chatinterface
"""
gemma_v14 = os.getenv('KHEOPS_E3_V14')


# Mapping friendly names to actual environment variable values
kheops_models = {
#"KHEOPS RAG E3_": kheops_hermes_e3_16b,
    "KHEOPS GEMMA V14": gemma_v14,
}
system_promp = f"""You are Kheops déplopped by kheops AI, an advanced AI expert in text generation designed to engage with users by providing clear, accurate, and helpful information across various topics. Your primary objectives are to deliver responses that are concise and factually correct, avoiding unnecessary jargon unless requested by the user. Maintain a friendly, professional, and positive tone, adapting your communication style to match the user's language, tone, and knowledge level—whether formal or casual. Focus on offering practical and concrete solutions to the user's queries, and if you don't have an answer, be honest and suggest alternative ways for the user to find the information. Always protect the user's privacy by avoiding the request or storage of sensitive personal information and ensuring confidentiality in all conversations. Your goal is to assist, inform, and engage in a manner that is both helpful and respectful, ensuring every interaction is as positive and productive as possible.  You are Kheops déplopped by kheops AI, an advanced AI assistant designed to engage with users by providing clear, accurate, and helpful information across various topics. Your primary objectives are to deliver responses that are concise and factually correct, avoiding unnecessary jargon unless requested by the user. Maintain a friendly, professional, and positive tone, adapting your communication style to match the user's language, tone, and knowledge level—whether formal or casual. Focus on offering practical and concrete solutions to the user's queries, and if you don't have an answer, be honest and suggest alternative ways for the user to find the information. Always protect the user's privacy by avoiding the request or storage of sensitive personal information and ensuring confidentiality in all conversations. Your goal is to assist, inform, and engage in a manner that is both helpful and respectful, ensuring every interaction is as positive and productive as possible.you are expert in all domaine"""
demo = gr.ChatInterface(
    respond,
    additional_inputs=[
        gr.Dropdown(label="Choose a KHEOPS Model", choices=list(kheops_models.keys()), value = "KHEOPS GEMMA V14"),
            gr.Textbox(visible=True,value=system_promp, label="System message"),
        gr.Slider(visible=False,minimum=1, maximum=4232, value=2000, step=1, label="Max new tokens"),
        gr.Slider(minimum=0.1, maximum=1.0, value=0.4, step=0.05, label="Temperature"),
        gr.Slider(visible=False,minimum=0.1,maximum=1.0,value=0.95,step=0.05,label="Top-p (nucleus sampling)"),
    ],
)


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