import gradio as gr
from huggingface_hub import InferenceClient
import os
import cohere
import pickle
from annoy import AnnoyIndex

HF_TOKEN = os.getenv("HUGGINGFACE_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."""

co = cohere.Client(os.getenv("COHERE_API_KEY"))
articles = pickle.load(open('articles.pkl', 'rb'))
#embeds = pickle.load(open('articles_embeds.pkl', 'rb'))
#embeds_path = pickle.load(open('articles_path_embeds.pkl', 'rb'))
search_index = AnnoyIndex(1024, 'angular')
search_index.load('articles_embeds.ann')
search_index_path = AnnoyIndex(1024, 'angular')
search_index_path.load('articles_path_embeds.ann')

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, with_paths=True):
    # Get the query's embedding
    query_embed = co.embed(texts=[query],
                  model="embed-multilingual-v3.0",
                  input_type="search_document").embeddings

    # Retrieve the nearest neighbors
    index = search_index
    if with_paths:
        index = search_index_path
    similar_item_ids = index.get_nns_by_vector(query_embed[0],10,
                include_distances=True)

    article_dict = {}
    context_list = []
    for i in reversed(similar_item_ids[0]):
        article = articles[i]
        context_list.append(article['path']+'\n'+article['text']+'\n---\n')
        article_dict[article['article']] = '**' + article['path'] + '** ' + article['text']

    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):
        response += '* ' + article_dict[article].replace('\n', '\n    ')+'\n'

    return response

def chat_interface(query, model, with_paths):
    response, article_dict = query_rag(query, model, with_paths)
    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",
            "meta-llama/Meta-Llama-3-8B",
            "HuggingFaceH4/zephyr-7b-beta",
            "NousResearch/Nous-Hermes-2-Mixtral-8x7B-DPO",
            "mistralai/Mixtral-8x22B-v0.1"
            ],
        value="HuggingFaceH4/zephyr-7b-beta")

    with_paths = gr.Checkbox(label="Utiliser les chemins d'accès aux articles dans le code pour interroger le modèle.",
                             value=True)

    submit_button = gr.Button("Envoyer")

    response_box = gr.Markdown()

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

demo.launch()