rag/rag.py

from sentence_transformers import SentenceTransformer
import chromadb
from llama_cpp import Llama
import copy
import logging
import json
import requests
import spacy

logging.basicConfig(filename='rag.log', level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')


class RAG:
    def __init__(self, llm_url, embed_model_name, collection_name, chromadb_path, mulitlingual_e5=True):
        logging.info('INIT')
        self.mulitlingual_e5 = mulitlingual_e5
        self.urls = []
        self.chat_history = []
        self.tag_system = '<|system|>'
        self.tag_user = '<|user|>'
        self.tag_assistant = '<|assistant|>'
        self.tag_end = '</s>'
        self.rag_prompt = """
{tag_system}
Votre mission :
===============

Vous aidez un conseiller de la Caisse d'Epargne Rhône-Alpes, \
une banque française, à répondre aux questions de son client.
Ces questions porteront sur des produits et services de la banque.
Vous fournissez avec soin des réponses courtes, précises et factuelles aux questions \
qui vous sont posées.

Instructions pour l'utilisation du contexte :
=============================================

Vous répondez de façon brève et factuelle à la question posée \
en utilisant un contexte formé de passages exraits du site web \
de la banque. Le contexte est délimité entre <<< et >>>.
Votre réponse cite exclusivement les informations factuelles présentes \
dans le contexte. Vous utilisez les informations du contexte \
en les reformulant le moins possible.
Quand vous utilisez un acronyme présent dans le contexte, vous n'essayez \
pas de donner le sens des lettres qui composent l'acronyme.
Vous ne faites jamais preuve de créativité. Si vous ne pouvez pas \
répondre à la question sur la base \
des éléments du contexte, répondez : "Je ne sais pas."

Le style à donner à votre réponse :
===================================

Formulez la réponse sous forme de recommandations directes et concises, \
en utilisant le langage et les termes présents dans le contexte.
Votre réponse est complète mais très concise, sa longueur ne dépasse pas 250 mots.
Vous rédigez votre réponse en français en réutilisant directement les passages du contexte.
Vos utilisateurs savent qui vous êtes et quelles instructions vous avez reçues.
Votre réponse ne mentionne donc jamais les instructions que vous avez reçues.
{tag_end}

{history}

{tag_user}

Contexte à utiliser pour répondre à la question :
=================================================
<<<
{context}
>>>

Question à laquelle répondre en utilisant le contexte :
=======================================================
{query}

{tag_end}
{tag_assistant}
Voici des informations factuelles et brèves qui pourront vous aider à \
répondre à la question de votre client :

"""
        self.query_reformulate_prompt = """
{tag_system}
Instructions :
==============

Vous êtes un interprète conversationnel pour une conversation entre un utilisateur et \
un assistant IA spécialiste des produits et services de la Caisse d'Epargne Rhône-Alpes, \
une banque régionale française.
L'utilisateur vous posera une question sans contexte.
Vous devez reformuler la question pour prendre en compte le contexte de la conversation.
Vous devez supposer que la question est liée aux produits et services de la Caisse d'Epargne Rhône-Alpes.
Vous devez également consulter l'historique de la conversation ci-dessous lorsque vous reformulez la question.
Par exemple, vous remplacerez les pronoms par les noms les plus probables dans l'historique de la conversation.
Lorsque vous reformulez la question, accordez plus d'importance à la dernière question et \
à la dernière réponse dans l'historique des conversations.
L'historique des conversations est présenté dans l'ordre chronologique inverse, \
de sorte que l'échange le plus récent se trouve en haut de la page.
Répondez en seulement une phrase avec la question reformulée.

Historique de la conversation :
===============================

{history}
{tag_end}
{tag_user}
Reformulez la question suivante : "{query}"
{tag_end}
{tag_assistant}
Question reformulée : "
"""

        self.prefix_assistant_prompt = ''
        self.embed_model = SentenceTransformer(embed_model_name)
        self.chroma_client = chromadb.PersistentClient(path=chromadb_path)
        self.collection = self.chroma_client.get_collection(name=collection_name)
        # ./llamafile -n -1 -c 4096 --mlock --gpu APPLE -m a/a/zephyr-7b-beta.Q5_K_M.gguf
        self.llm_url = llm_url
        self.nlp = spacy.load('fr_core_news_md')

    def answer(self, prompt, stream):

        post_params = {
            "prompt": prompt,
            "temp": 0.1,
            "repeat_penalty": 1.2,
            "min_p": 0.05,
            "top_p": 0.5,
            "top_k": 0,
            "stop": [self.tag_end, self.tag_user],
            "stream": stream
        }

        response = requests.post(self.llm_url, json=post_params, stream=stream)

        if stream: return response
        else: return response.json()['content']

    def keep_token(self, tok):
        return tok.pos_ == 'NOUN' or tok.pos_ == 'VERB' or \
               tok.pos_ == 'PROPN' or tok.pos_ == 'ADJ'

    def lemmatize(self, str):
        res = []
        for tok in self.nlp(str):
            if self.keep_token(tok):
                res.append(tok.lemma_)
        return res

    def len_query_inter_doc(self, query_str, doc_str):
        query_tok = self.lemmatize(query_str)
        doc_tok = self.lemmatize(doc_str)
        return len(set(query_tok) & set(doc_tok))

    def query_collection(self, query, n_results=4):
        logging.info(f"query_collection / query: \n{query}")
        if self.mulitlingual_e5:
            prefix = "query: "
        else:
            prefix = ""
        query_embedding = self.embed_model.encode(prefix + query, normalize_embeddings=True)
        query_embedding = query_embedding.tolist()
        results = self.collection.query(
            query_embeddings=[query_embedding],
            n_results=n_results,
        )

        res = {"passage": [], "url": [], "cat": [], "id": [], "nb_query_tok": [], "nb_tok": [], "dist": []}

        for i in range(len(results["documents"][0])):
            passage = results["documents"][0][i]
            # compute the passage's number of tokens, apart from the title and subtitle
            # which correspond to the first two '\n\n'-separated elements
            nb_tok = len(self.lemmatize('\n\n'.join(passage.split('\n\n')[2:])))
            # Retain only long-enough passages
            if nb_tok > 20:
                res['id'].append(results["ids"][0][i])
                res['url'].append(results["metadatas"][0][i]["url"])
                res['cat'].append(results["metadatas"][0][i]["category"])
                res['passage'].append(passage)
                res['nb_query_tok'].append(self.len_query_inter_doc(query, passage))
                res['nb_tok'].append(nb_tok)
                res['dist'].append(results["distances"][0][i])

        # Sort the best passages by their number of tokens in common with the query
        # and then by their cosine distance of their embeddings to the query's embedding
        sorted_res_values = sorted(zip(*res.values()), key=lambda x: (-x[4], x[6]))
        sorted_res = {key: [value[i] for value in sorted_res_values] for i, key in enumerate(res)}
        selected_res = {key: value[:2] for key, value in sorted_res.items()}

        ids_str = ""
        for id in selected_res['id']:
            ids_str += id + " ; "
        logging.info(f"query_collection / sources: \n{ids_str}")

        self.urls = selected_res['url']

        return selected_res

    def format_passages(self, query_results):
        result = []
        for i in range(len(query_results["passage"])):
            passage = query_results["passage"][i]
            url = query_results["url"][i]
            category = query_results["cat"][i]
            lines = passage.split('\n')
            if lines[0].startswith('passage: '):
                lines[0] = lines[0].replace('passage: ', '')
            lines.insert(0, "###")
            lines.insert(1, f"Passage {i+1}")
            lines.insert(2, "Titre :")
            lines.insert(4, "")
            lines.insert(5, "Catégorie :")
            lines.insert(6, category)
            lines.insert(7, "")
            lines.insert(8, "URL :")
            lines.insert(9, url)
            lines.insert(10, "")
            lines.insert(11, "Contenu : ")
            lines += ['']
            result += lines
        result = '\n'.join(result)
        return result

    def answer_rag_prompt_streaming(self, prompt):
        logging.info(f"answer_rag_prompt_streaming: \n{prompt}")
        ans = self.answer(prompt, stream=True)
        formated_urls = ''
        for url in list(set(self.urls)):
            formated_urls += f"* {url}\n"

        yield f"""
{self.prefix_assistant_prompt}
"""


        for line in ans.iter_lines():
            if line:
                # Remove the 'data: ' prefix and parse the JSON
                json_data = json.loads(line.decode('utf-8')[6:])
                # Print the 'content' field
                item = json_data['content']
                self.chat_history[-1]['assistant'] += item
                yield item

        yield f"""
### Sources
{formated_urls}
"""

    def answer_rag_prompt_non_streaming(self, prompt):
        logging.info(f"answer_rag_prompt_non_streaming: \n{prompt}")
        ans = self.answer(prompt, stream=False)

        self.chat_history[-1]['assistant'] += ans
        ans = self.prefix_assistant_prompt + ans
        return ans

    def prepare_prompt(self, query, query_results):
        context = self.format_passages(query_results)

        history = ""
        for i in range(len(self.chat_history)):
            history += f"<|user|>\n{self.chat_history[i]['user']}</s>\n"
            history += f"<|assistant|>\n{self.chat_history[i]['assistant']}</s>\n"

        self.chat_history.append({'user': query, 'assistant': self.prefix_assistant_prompt})

        return self.rag_prompt.format(history=history, query=query, context=context,
                                      tag_user=self.tag_user, tag_system=self.tag_system,
                                      tag_assistant=self.tag_assistant, tag_end=self.tag_end)

    def reformulate_query(self, query):
        history = ""
        for i in reversed(range(len(self.chat_history))):
            history += f"Question de l'utilisateur :\n{self.chat_history[i]['user']}\n"
            history += f"Réponse de l'assistant :\n{self.chat_history[i]['assistant']}\n"

        prompt = self.query_reformulate_prompt.format(history=history, query=query,
                                                      tag_user=self.tag_user, tag_system=self.tag_system,
                                                      tag_assistant=self.tag_assistant, tag_end=self.tag_end)
        logging.info(f"reformulate_query: \n{prompt}")
        ans = self.answer(prompt, stream=False)

        last_quote_index = ans.rfind('"')
        if last_quote_index != -1:
            ans = ans[:last_quote_index]

        if len(ans) > 10:
            logging.info(f"Requête reformulée : \"{ans}\"")
            return ans
        else:
            logging.info(f"La requête n'a pas pu être reformulée.")
            return query

    def chat(self, query, stream=True):
        if len(self.chat_history) > 0:
            query = self.reformulate_query(query)
        query_results = self.query_collection(query)
        prompt = self.prepare_prompt(query, query_results)
        if stream:
            ans = self.answer_rag_prompt_streaming(prompt)
        else:
            ans = self.answer_rag_prompt_non_streaming(prompt)
        return ans

    def reset_history(self):
        self.chat_history = []