Cross encodeur sur la sortie du RAG

Directement dans le tool

AgentReact/start.py

@@ -22,4 +22,4 @@ initial_input = {
 config={"configurable": {"thread_id": 'yes'}}
 
 # Et je lance !
-streamGraph(initial_input, config, getGraph(), showSysMessages=True)
+streamGraph(initial_input, config, getGraph(), showSysMessages=False)

AgentReact/utils/state.py

@@ -6,9 +6,6 @@ import operator
 class CustomState(MessagesState):
     todo: Annotated[list, operator.add]  # Les tâches en cours, au format JSON
 
-    ragQuery: str # Requête envoyée au RAG, pour le cross-encodeur
-    ragDocuments: List[str] # Documents retrouvés par le RAG, pour le cross-encodeur
-
     lastSummarizedMessage: int # Index du message où l'on s'était arrêté de résumer
 
     stop: bool # Permet d'indiquer la fin de l'exécution de l'agent

AgentReact/utils/tools.py

@@ -8,11 +8,21 @@ from pathlib import Path
 from typing import List, Dict, Annotated, Tuple
 import sys
 import os
+from sentence_transformers import CrossEncoder
 from langgraph.types import interrupt
+from langchain_core.documents import Document
 
 from .StateElements.TodoElement import TodoElement
 from .VectorDatabase import VectorDatabase
 from .InterruptPayload import InterruptPayload
+from langchain_mistralai import ChatMistralAI # LLM définit dans le fichier agent
+
+CROSS_ENCODEUR = CrossEncoder('jinaai/jina-reranker-v2-base-multilingual', trust_remote_code=True)
+CROSS_ENCODEUR_MIN_SIM_SCORE = 0.5
+
+llm = ChatMistralAI(model="ministral-3b-2512", # Petit modèle, pour aller vite sur des tâches simples
+    temperature=0,
+    max_retries=2)
 
 @tool
 def append_part_to_report(contenu:str)->str:
@@ -219,15 +229,16 @@ def search_in_files(query:str, state: Annotated[dict, InjectedState])->str:
 
     retrieved_docs = bdd.similarity_search(query, k=5) # 5 documents
 
+    reprompt = cross_encodeur(query, retrieved_docs) # Cross-encodeur en charge de regarder si la recherche a été efficace ou non
+    if reprompt is not None:
+        # Il y a un nouveau prompt, il faut recommencer la recherche
+        retrieved_docs = bdd.similarity_search(reprompt, k=5)
+
     # Conversion des documents en texte
     docs_content = "\n".join(
         [f"Document {i+1}:\n{doc.page_content}" for i,doc in enumerate(retrieved_docs)]
     )
 
-    # Sauvegarde des données dans le State
-    state["ragQuery"] = query
-    state["ragDocuments"] = retrieved_docs
-
     return docs_content # Retourne la liste de documents trouvés
 
 @tool
@@ -298,4 +309,32 @@ def getWeeklyReportTools()->List['Tools']:
     """
     Récupérer la liste des tools, POUR LE LLM EN CHARGE DE FAIRE LES RAPPORTS DE CHAQUE SEMAINE
     """
-    return [write_week_report, write_library_tools_details_on_internship, internet_search, search_in_files]
+    return [write_week_report, write_library_tools_details_on_internship, internet_search, search_in_files]
+
+
+
+# CROSS-ENCODEUR
+# Selon https://app.ailog.fr/fr/blog/guides/cross-encoder-reranking
+def cross_encodeur(query:str, docs:List[Document])->str|None:
+    """
+    Fonction que j'utilise pour faire tourner le cross-encodeur.
+    Il vérifie la sortie de la recherche des documents, et reformule la question si besoin.
+
+    Args:
+        query (str): Requête originale
+        docs (List[Document]): documents retrouvés par la première recherche
+
+    Returns:
+        str|None: None si le résultat est valide, une reformulation de la requête sinon.
+    """
+    pairs = [[query, doc.page_content] for doc in docs]
+    scores = CROSS_ENCODEUR.predict(pairs) # Scores de similarité de la recherche
+    
+    sum = 0
+    for s in scores: sum+= s
+    if sum / len(docs) < CROSS_ENCODEUR_MIN_SIM_SCORE:
+        # Si en dessous d'un certain score de qualité
+        return llm.invoke("Voici une recherche de documents locale.\
+            Essaie de la réecrire de façon à améliorer le score de la recherche.\
+            Ne retourne QUE la nouvelle question, rien d'autre. Voici la question originale: "+query).content
+    return None # Recherche valide, pas besoin de retenter