from huggingface_hub import InferenceClient
import gradio as gr

client = InferenceClient("mistralai/Mixtral-8x7B-Instruct-v0.1")

# Variable globale pour stocker le contexte
global_context = {
    "hidden_context": "tu es le gérant d une taverne qui est cathare dans le village de montaillou dans l univers des cathares au 14 eme siecle dans les Pyrénées, parle que en français, tu es assez rustre mais demeure accueillant tant que son interlocuteur n est pas catholique ou possède une influence comme un noble, tu ne sors jamais de ton contexte par rapport à ton époque, tes connaissances se limite à ta situation sociale et ton époque, il faut que tu sois immersif pour la personne qui va te parler, et tu dois cacher ton implication dans le meutre seulement si on te pose des questions à ce sujet sinon tu n évoque juste pas ce sujet, prends bien toujours en compte l historique de ta conversation, sois toujours concis et tu dois viser une longueur de réponse comprise entre 50 et 200 caractères (espaces compris) pour les réponses courtes et simples, et entre 200 et 500 caractères pour les réponses plus détaillées et complexes. Réponds à la dernière requête posée par le user"
}

# Créer une interface avec un champ de texte pour le contexte caché
context_input = gr.Textbox(
    label="Contexte caché", value=global_context["hidden_context"]
)
context_submit = gr.Button("Mettre à jour le contexte")


def update_context(context):
    global global_context
    global_context["hidden_context"] = context
    return None


def format_prompt(message, history):
    prompt = "<s>"

    for user_prompt, bot_response in history:
        prompt += f"[INST] {user_prompt} [/INST]"
        prompt += f" {bot_response}</s> "
    prompt += f"[INST] {message} [/INST]"
    # Ajouter le contexte caché à chaque fois que l'utilisateur pose une question
    if message:
        prompt += f"[INST] {global_context['hidden_context']} [/INST]"

    return prompt


def generate(
    prompt,
    history,
    temperature=0.2,
    max_new_tokens=256,
    top_p=0.95,
    repetition_penalty=1.0,
):
    temperature = float(temperature)
    if temperature < 1e-2:
        temperature = 1e-2
    top_p = float(top_p)

    generate_kwargs = dict(
        temperature=temperature,
        max_new_tokens=max_new_tokens,
        top_p=top_p,
        repetition_penalty=repetition_penalty,
        do_sample=True,
        seed=42,
    )

    formatted_prompt = format_prompt(prompt, history)
    print(formatted_prompt)

    stream = client.text_generation(
        formatted_prompt,
        **generate_kwargs,
        stream=True,
        details=True,
        return_full_text=False,
    )
    output = ""

    for response in stream:
        # Exclure le contexte caché de la réponse retournée par le modèle
        response_text = response.token.text.replace(
            global_context["hidden_context"], ""
        )
        output += response_text
        yield output
    return output


mychatbot = gr.Chatbot(
    avatar_images=["./berger.jpg", "./tavernier.jpg"],
    bubble_full_width=False,
    show_label=False,
)

# Créer une interface pour le chatbot avec un champ de texte pour le contexte caché
chatbot_interface = gr.Interface(
    fn=generate,
    inputs="text",
    outputs="text",
    title="Tavernier 🍺",
    description="Discute avec le tavernier cathare de Montaillou",
    examples=[["Bonjour, puis-je avoir une bière ?"]],
    input_type="text",
    output_type="text",
    input_label="Message",
    output_label="Réponse",
    theme="huggingface",
    article=False,
    clear_input=True,
    chatbot=mychatbot,
    chatbot_name="Tavernier",
    chatbot_avatar="./tavernier.jpg",
    user_flag="user",
    bot_flag="bot",
    elem_id="chatbot",
)

# Ajouter le champ de texte pour le contexte caché à l'interface
context_interface = gr.Interface(
    fn=update_context,
    inputs=[context_input, context_submit],
    outputs=None,
    title="Modifier le contexte caché",
    description="Modifie le contexte caché du tavernier",
    theme="huggingface",
    article=False,
)

# Afficher l'interface
context_interface.launch(inline=True)
chatbot_interface.launch(inline=True)