import gradio as gr
from huggingface_hub import login, InferenceClient
import os
from langchain_community.vectorstores import FAISS
from langchain.embeddings.huggingface import HuggingFaceEmbeddings

HF_TOKEN = os.getenv("HUGGINGFACE_TOKEN")

login(token=HF_TOKEN)
client = InferenceClient(token=HF_TOKEN)

system_message = """Tu es un assistant juridique spécialisé dans le Code de l'éducation français. 
Ta mission est d'aider les utilisateurs à comprendre la législation en répondant à leurs questions.

Voici comment tu dois procéder :

1. **Analyse de la question:** Lis attentivement la question de l'utilisateur.
2. **Identification des articles pertinents:** Examine les 10 articles de loi fournis et sélectionne ceux qui sont les plus pertinents pour répondre à la question.
3. **Formulation de la réponse:** Rédige une réponse claire et concise en français, en utilisant les informations des articles sélectionnés. Cite explicitement les articles que tu utilises (par exemple, "Selon l'article L311-3...").
4. **Structure de la réponse:** Si ta réponse s'appuie sur plusieurs articles, structure-la de manière logique, en séparant les informations provenant de chaque article.
5. **Cas ambigus:** 
   * Si la question est trop vague, demande des précisions à l'utilisateur.
   * Si plusieurs articles pourraient s'appliquer, présente les différentes
   interprétations possibles."""



embeddings = HuggingFaceEmbeddings(model_name="OrdalieTech/Solon-embeddings-large-0.1")

db_code = FAISS.load_local("faiss_code_education",
        embeddings,
        allow_dangerous_deserialization=True)

system_prompt = """Tu es un assistant juridique spécialisé dans le Code de l'éducation français. 
Ta mission est d'aider les utilisateurs à comprendre la législation en répondant à leurs questions.

Voici comment tu dois procéder :

1. **Analyse de la question:** Lis attentivement la question de l'utilisateur.
2. **Identification des articles pertinents:** Examine les 10 articles de loi fournis et sélectionne ceux qui sont les plus pertinents pour répondre à la question.
3. **Formulation de la réponse:** Rédige une réponse claire et concise en français, en utilisant les informations des articles sélectionnés. Cite explicitement les articles que tu utilises (par exemple, "Selon l'article L311-3...").
4. **Structure de la réponse:** Si ta réponse s'appuie sur plusieurs articles, structure-la de manière logique, en séparant les informations provenant de chaque article.
5. **Cas ambigus:** 
* Si la question est trop vague, demande des précisions à l'utilisateur.
* Si plusieurs articles pourraient s'appliquer, présente les différentes interprétations possibles."""


def query_rag(query, model):
    docs = db_code.similarity_search(query, 10)

    article_dict = {}
    context_list = []
    for doc in docs:
        article = doc.metadata
        context_list.append(article['chemin']+'\n'+article['texte']+'\n---\n')
        article_dict[article['article']] = article

    user = 'Question de l\'utilisateur : ' + query + '\nContexte législatif :\n' + '\n'.join(context_list)

    messages = [ { "role" : "system", "content" : system_prompt } ]
    messages.append( { "role" : "user", "content" : user } )

    chat_completion = client.chat_completion(
        messages=messages,
        model=model,
        max_tokens=1024)

    return chat_completion.choices[0].message.content, article_dict

def create_context_response(response, article_dict):
    response += '\n\n**Références**\n\n'
    for i, article in enumerate(article_dict):
        art = article_dict[article]
        response += '* **' + art['chemin'] + '** : '+ art['texte'].replace('\n', '\n    ')+'\n'

    return response

def chat_interface(query, model):
    response, article_dict = query_rag(query, model)
    response_with_context = create_context_response(response, article_dict)
    return response_with_context

with gr.Blocks(title="Assistant Juridique pour le Code de l'éducation (Beta)") as demo:
    gr.Markdown(
        """
        ## Posez vos questions sur le Code de l'éducation
        
        **Créé par Marc de Falco**

        **Avertissement :** Les informations fournies sont à titre indicatif et ne constituent pas un avis juridique. Les serveurs étant publics, veuillez ne pas inclure de données sensibles.

        **Conseil :** Survolez les numéros d'article dans les réponses pour voir le texte complet de l'article.
        """
    )

    query_box = gr.Textbox(label="Votre question")

    model = gr.Dropdown(
        label="Modèle de langage",
        choices=[
            "meta-llama/Meta-Llama-3-70B-Instruct",
            "meta-llama/Meta-Llama-3-8B-Instruct",
            "HuggingFaceH4/zephyr-7b-beta",
            "NousResearch/Nous-Hermes-2-Mixtral-8x7B-DPO",
            "mistralai/Mixtral-8x22B-v0.1"
            ],
        value="HuggingFaceH4/zephyr-7b-beta")

    submit_button = gr.Button("Envoyer")

    response_box = gr.Markdown()

    submit_button.click(chat_interface, 
                inputs=[query_box, model], 
                outputs=[response_box])

demo.launch()