core.py

import os
os.environ["REQUESTS_CA_BUNDLE"] = ""

from llama_index.core import (
    VectorStoreIndex,
)

import chromadb
from llama_index.llms.mistralai import MistralAI
from llama_index.embeddings.mistralai import MistralAIEmbedding
from llama_index.vector_stores.chroma import ChromaVectorStore
from llama_index.core.storage.storage_context import StorageContext
from llama_index.core import ServiceContext
from llama_index.core.retrievers import VectorIndexRetriever

import tenacity
from typing import List, Union
import textwrap
import dsp
import dspy
from dsp.utils import dotdict
from dsp.modules.lm import LM


MISTRAL_API_KEY = "Yb2kAF0DR4Mva5AEmoYFV3kYRAKdXB7i"
EMBEDDING_MODEL_NAME = "mistral-embed"
COMPLETION_MODEL_NAME = "mistral-medium"
PERSISTANT_DB_PATH = "./chroma_db"

def deduplicate_with_sources(seq: list[str], sources: list[str]) -> list[str]:
    final_context = []
    final_sources = []

    for snippet, source in zip(seq, sources):
        if snippet not in final_context:
            final_context.append(snippet)
            final_sources.append(source)

    return final_context, final_sources


class FinalAnswerChainOfThought(dspy.Predict):
    def __init__(self, signature, rationale_type=None, activated=True, **config):
        super().__init__(signature, **config)

        self.activated = activated

        signature = self.signature
        *keys, last_key = signature.kwargs.keys()

        DEFAULT_RATIONALE_TYPE = dsp.Type(
            prefix="Raisonnement: Décompose la requête en sous requêtes, réponds à chacune d'entre elle d'abord avant de",
            desc="produire la ${" + last_key + "}. Nous ...",
        )

        rationale_type = rationale_type or DEFAULT_RATIONALE_TYPE

        extended_kwargs = {key: signature.kwargs[key] for key in keys}
        extended_kwargs.update(
            {"rationale": rationale_type, last_key: signature.kwargs[last_key]}
        )

        self.extended_signature = dsp.Template(
            signature.instructions, **extended_kwargs
        )

    def forward(self, **kwargs):
        new_signature = kwargs.pop("new_signature", None)
        if new_signature is None:
            if self.activated is True or (
                self.activated is None and isinstance(dsp.settings.lm, dsp.GPT3)
            ):
                signature = self.extended_signature
            else:
                signature = self.signature
        else:
            signature = dsp.Template(self.signature.instructions, **new_signature)
        return super().forward(signature=signature, **kwargs)


    def dump_state(self):
        state = super().dump_state()

        # Cache the signature instructions and the last field's name.
        state["extended_signature_instructions"] = self.extended_signature.instructions
        state["extended_signature_prefix"] = self.extended_signature.fields[-1].name

        return state

    def load_state(self, state):
        super().load_state(state)
        
        # Reconstruct the signature.
        if "extended_signature_instructions" in state:
            instructions = state["extended_signature_instructions"]
            self.extended_signature.instructions = instructions
        
        if "extended_signature_prefix" in state:
            prefix = state["extended_signature_prefix"]
            self.extended_signature.fields[-1] = self.extended_signature.fields[-1]._replace(name=prefix)


class GenerateSearchQuery(dspy.Signature):
    """Écris en une phrase une requête de recherche simple qui aidera à répondre à une question complexe. Utilise des termes plus génériques pour faciliter la recherche"""

    contexte = dspy.InputField(desc="peut contenir des faits pertinents")
    question = dspy.InputField()
    requete = dspy.OutputField()


class GenerateAnswer(dspy.Signature):
    """Fournis une réponse concise et factuelle en faisant référence au contexte"""

    contexte = dspy.InputField(desc="peut contenir des faits pertinents")
    question = dspy.InputField()
    reponse = dspy.OutputField(desc="contient obligatoirement des références du contexte ([numéro de référence dans le contexte])")


class SimplifiedBaleen(dspy.Module):
    def __init__(self, passages_per_hop=5, max_hops=2):
        super().__init__()

        self.generate_query = [dspy.ChainOfThought(GenerateSearchQuery) for _ in range(max_hops)]
        self.retrieve = dspy.Retrieve(k=passages_per_hop)
        self.generate_answer = FinalAnswerChainOfThought(GenerateAnswer)
        self.max_hops = max_hops
    
    @tenacity.retry(
            stop=tenacity.stop_after_attempt(3),
            wait=tenacity.wait_fixed(10)
    )
    def forward(self, question):
        context = []
        sources = []

        for hop in range(self.max_hops):
            query = self.generate_query[hop](contexte=context, question=question).requete
            retriever_output = self.retrieve(query)
            passages, current_sources = retriever_output
            passages = passages.passages
            current_sources = current_sources.sources
            context, sources = deduplicate_with_sources(context + passages, sources + current_sources)
        print("\nFinal context:")
        for idx, (elem, source) in enumerate(zip(context, sources)):
            print(f"[{idx + 1}] : {elem} ==> {source}")
        pred = self.generate_answer(contexte=context, question=question)
        return dspy.Prediction(contexte=context, answer=pred.reponse), context, sources
    
class ChromaDBRetriever(dspy.Retrieve):
    def __init__(self, vector_store_index:VectorStoreIndex, *args, **kwargs):
        self.vector_store_index = vector_store_index
        # self.similarity_top_k = kwargs.get("k", 2)
    
    def forward(self, query_or_queries: Union[str, List[str]], *args, **kwargs) -> dspy.Prediction:
        # query_engine = self.vector_store_index.as_retriever(similarity_top_k=kwargs.get("k"), vector_store_query_mode="text_search")

        query_engine = VectorIndexRetriever(
            index=self.vector_store_index,
            similarity_top_k=kwargs.get("k"),
        )

        top_k = query_engine.retrieve(query_or_queries)# k=kwargs.get("k", 1))
        passages = [dotdict({"long_text": doc.text, "source_name": doc.metadata["source_name"]}) for doc in top_k]

        return passages
    
class DspyMistralWrapper(LM):
    def __init__(self, llm_model: MistralAI):
        super().__init__(model="mistral")
        self.llm_model = llm_model
        self.history = []

    def basic_request(self, prompt, **kwargs):
        response = self.llm_model.complete(prompt)
        # res_data = response
        history = {
            "prompt": prompt,
            "response": response.text,
            "kwargs": kwargs,
            "raw_kwargs": {},
        }
        self.history.append(history)
        return response.text
    
    def __call__(self, prompt, only_completed=True, return_sorted=False, **kwargs):
        return [self.basic_request(prompt=prompt)]

def init_llm(persistant_db_path=PERSISTANT_DB_PATH):
    llm = MistralAI(api_key=MISTRAL_API_KEY, model=COMPLETION_MODEL_NAME, temperature=0, max_tokens=10000)
    embed_model = MistralAIEmbedding(model_name=EMBEDDING_MODEL_NAME, api_key=MISTRAL_API_KEY)

    db = chromadb.PersistentClient(path=persistant_db_path)
    chroma_collection = db.get_or_create_collection("quickstart")

    # set up ChromaVectorStore and load in data
    vector_store = ChromaVectorStore(chroma_collection=chroma_collection)
    storage_context = StorageContext.from_defaults(vector_store=vector_store)
    service_context = ServiceContext.from_defaults(
        chunk_size=1024, llm=llm, embed_model=embed_model
    )
    index = VectorStoreIndex(
        [], service_context=service_context, storage_context=storage_context
    )

    retriever = ChromaDBRetriever(vector_store_index=index)

    dspy_llm = DspyMistralWrapper(llm_model=llm)
    dspy.settings.configure(lm=dspy_llm, rm=retriever)
    dspy.settings.configure(lm=dspy_llm, rm=retriever)

def invoke(input_prompt:str):
    uncompiled_baleen = SimplifiedBaleen(passages_per_hop=5, max_hops=2)
    predictions, justifications, sources = uncompiled_baleen(input_prompt)
    content = " ".join(textwrap.wrap(predictions.answer))
    # import time
    # time.sleep(0.2)
    # content = "Le numéro d'identification d'un bovin en France est composé de 12 chiffres. Les deux premiers chiffres représentent le code pays "FR" ([1]). Les deux chiffres suivants correspondent au numéro de code INSEE du département où se trouve l'animal au moment de son identification ([2]). Les huit chiffres restants sont attribués sous la responsabilité de l'établissement de l'élevage ([2]). Pour les bovins identifiés en France après le 1er septembre 1998, le numéro de travail est composé des 4 derniers chiffres du numéro national ([1])."
    # justifications = [
    #     "Page 20 sur 47  IDENTIFICATION Vv. BOVINE  Annexe de l’arrêté du 6 août 2013 relatif à l’identification des animaux de l’espèce bovine  Notification des informations  10. Notifier chaque événement dans les sept jours, à (nom du maître d'oeuvre de l'identification) : – soit en transmettant l’exemplaire du « DOCUMENT DE NOTIFICATION – REGISTRE BOVIN » prévu à cet effet ; – soit en transmettant ces informations par voie électronique, selon les modalités techniques définies par (nom du maître d'oeuvre de l'identification).  11. Notifier toute anomalie constatée sur tout document (passeport, livre des bovins, document de notification….) à (nom du maître d'oeuvre de l'identification).  Pertes de marques auriculaires agréées numérotées  12. En cas de perte, de détérioration ou d’illisibilité d'une seule marque auriculaire agréée, commander à (nom du maître d'oeuvre de l'identification) une marque auriculaire agréée permettant d'avoir toujours le même numéro national sur ce bovin et l’apposer au plus vite, dans un délai maximum de trente jours après la livraison.  13. En cas de perte de détérioration ou d’illisibilité de deux marques auriculaires agréées, isoler l'animal et faire appel à (nom du maître d'oeuvre de l'identification) pour la vérification de l’identité de l'animal et le remplacement éventuel de ses marques auriculaires agréées à l'identique. En cas d'impossibilité de reconnaissance de l’identité de l'animal, ce dernier pourra être détruit sans compensation financière, conformément à la réglementation communautaire en vigueur.  Circulation des animaux  14. Ne laisser entrer dans mon exploitation un bovin, ou en sortir, que correctement identifié (deux marques auriculaires agréées numérotées, passeport correctement renseigné et correspondant aux caractéristiques de l'animal).  Cessation d’activité  15. Informer (nom du maître d'oeuvre de l'identification) de ma cessation d'activité.  Restitution du matériel d'identification  16. Restituer à (nom du maître d'oeuvre de l'identification), en cas de cessation d'activité, ou à sa demande, la totalité des marques d'identification et des passeports dont je dispose.  Dispositions générales  17. Sur demande d'un agent mandaté par l'établissement de l'élevage ou par (nom du maître d'oeuvre de l'identification) le cas échéant ou de tout agent mandaté de la direction départementale en charge de la protection de la population ou de la direction départementale des territoires, communiquer toute information utile et présenter tous mes animaux, toutes les marques auriculaires agréées en stock ainsi que tous les documents d'identification dont je dispose.  18. En cas d’intervention de ces agents, faciliter l'accès à mes animaux en assurant notamment leur contention.  19. Payer à (nom du maître d'oeuvre de l'identification) les sommes dont je suis redevable pour les opérations d'identification. En cas de non-paiement, (nom du maître d’oeuvre de l’identification) peut me refuser la délivrance des passeports.  20. En cas de non-respect de mes obligations, je dois avoir recours à un agent mandaté par (nom du maître d'oeuvre de l'identification) à mes frais, pour la réalisation de l'identification des animaux de mon exploitation.  21. Je suis informé que le non-respect de mes obligations peut se traduire par la perte des primes, voire l'obligation de paiement de pénalités financières complémentaires.  Date et signature.  Vu le détenteur.  un exemplaire signé est retourné à (nom du maître d'oeuvre de l'identification) un exemplaire est conservé par le détenteur",
    #     "Page 21 sur 47  IDENTIFICATION Vv. BOVINE  Annexe de l’arrêté du 6 août 2013 relatif à l’identification des animaux de l’espèce bovine  7.1.2  I-B. - Déclaration du détenteur-opérateur commercial auprès de l'établissement départemental/interdépartemental de l'élevage  *** Préciser dans le texte l'organisme qui assure la maîtrise d'oeuvre de l'identification***  Je soussigné, M. ..., détenteur, agissant en mon nom propre / agissant en tant que responsable de (nom de la personne morale)……………  déclare avoir pris connaissance de l'obligation qui m'est faite d'accomplir les opérations d'identification des bovins, détenus sur mon exploitation n° ……....., telles que prévues par la réglementation communautaire et nationale en vigueur.  Ces obligations portent plus particulièrement sur les points suivants :  Gestion des marques auriculaires agréées numérotées et des documents d'identification  1. Notifier à (nom du maître d'oeuvre de l'identification) toute perte de documents d'identification.  2. Ne déboucler sous aucun prétexte quelque animal que ce soit.  Tenue du registre  3. Inscrire sur le registre des bovins chaque événement : naissance, entrée, sortie ou mort. Remarque : dans le cas particulier d’une naissance, l’apposition des marques auriculaires sur l’animal, réalisée par l’EdE, s’effectue avant l’enregistrement de la naissance sur le registre. Le registre des bovins peut être tenu : – ou en utilisant les documents de notification mis à disposition par (nom du maître d'oeuvre de l'identification) ; – ou en utilisant le registre fiscal, – ou sur support électronique. Dans ce dernier cas, le registre doit pouvoir être présenté sur demande d’un agent identificateur désigné par l’EdE, d’un agent mandaté par la direction départementale de la protection de la population ou d’un agent mandaté par la direction départementale des territoires.  4. Vérifier que le registre des bovins contient l'ensemble des informations d'identification, tenues à jour, concernant mon exploitation.  5. Conserver dans le registre des bovins au minimum, les informations des trois dernières années.  Notification des informations  6. Notifier chaque événement dans les sept jours : – soit en transmettant à (nom du maître d'oeuvre de l'identification) un exemplaire du support papier du registre ; – soit en transmettant ces informations par voie électronique, selon les modalités techniques définies par (nom du maître d'oeuvre de l'identification).  7. Notifier toute anomalie constatée sur tout document (passeport, livre des bovins, document de notification….) à (nom du maître d'oeuvre de l'identification).  Pertes de marques auriculaires agréées numérotées  8. Notifier toute perte, détérioration ou illisibilité de marques auriculaires agréées à (nom du maître d'oeuvre de l'identification). Un agent identificateur habilité vérifiera l'identité de l'animal et effectuera le remplacement éventuel des marques auriculaires agréées à l'identique. En cas d'impossibilité de reconnaissance de l'identité de l'animal, ce dernier pourra être détruit sans compensation financière, conformément à la réglementation communautaire en vigueur.      z- | Liberd + Egalté » Fraternlté REPUBLIQUE FRANCAISE  Annexe v1.2"
    # ]
    # sources = [
    #     "Dossier-PAC-2022_notice_ICHN.pdfPAGENUMBER2",
    #     "lien2.html",
    # ]
    if not len(justifications) == len(sources):
        raise RuntimeError(f"Justifications ({len(justifications)}) and sources ({len(sources)}) does not have the same size")
    return {
        "content": content,
        "justifications": justifications,
        "sources": sources,
    }