Chat_Bot_Actuariat/tools.py

import os
import pandas as pd
import chromadb
from tqdm import tqdm
from dotenv import load_dotenv
from typing import List, Any # Ajout de Any pour corriger le warning

from langchain_community.tools import TavilySearchResults , DuckDuckGoSearchRun
from langchain_openai import ChatOpenAI
from langchain_chroma import Chroma
from langchain_huggingface import HuggingFaceEmbeddings
from langchain_core.documents import Document
from langchain_core.retrievers import BaseRetriever
from langchain_core.callbacks import CallbackManagerForRetrieverRun
from langchain_core.prompts import ChatPromptTemplate

# Correction des imports des chains (Chemins officiels 2026)
from langchain_classic.chains import create_retrieval_chain
from langchain_classic.chains.combine_documents import create_stuff_documents_chain

# Correction de l'import Mixedbread
from mixedbread_ai.client import MixedbreadAI

from langgraph.prebuilt import create_react_agent
from langgraph.checkpoint.memory import MemorySaver

load_dotenv()

# --- CONFIGURATION ---
CHROMA_PATH = "./chroma_db_actuariat"
PARQUET_PATH = "memoires_total_local.parquet"
MODEL_RRK = "mixedbread-ai/mxbai-rerank-large-v1"
LLM_NAME = "gpt-4o-mini"
EMB_MODEL_NAME = "mixedbread-ai/mxbai-embed-large-v1" 

mxbai_client = MixedbreadAI(api_key=os.environ.get("MXBAI_API_KEY"))

# --- 1. INITIALISATION DE LA BASE (CHROMA) ---
def get_vector_db():
    persistent_client = chromadb.PersistentClient(path=CHROMA_PATH)
    
    if "embeddings_mxbai" not in [c.name for c in persistent_client.list_collections()]:
        print("Initialisation de la base de données vectorielle locale...")
        collection = persistent_client.create_collection(name="embeddings_mxbai")
        
        # Chargement du Parquet
        df = pd.read_parquet(PARQUET_PATH)
        
        batch_size = 5000 
        for i in tqdm(range(0, len(df), batch_size), desc="Peuplement de ChromaDB"):
            batch = df.iloc[i : i + batch_size]
            
            # Sécurité : Conversion des ID en string
            ids = [str(idx) for idx in batch["id"].tolist()]
            
            collection.add(
                ids=ids,
                embeddings=batch["embedding"].tolist(), 
                documents=batch["document"].tolist(),
                metadatas=[{"title": str(m["title"])} for m in batch["metadata"]]
            )
        print("✅ Base de données prête.")
    
    return Chroma(
        client=persistent_client,
        collection_name="embeddings_mxbai",
        embedding_function=HuggingFaceEmbeddings(model_name=EMB_MODEL_NAME)
    )

# --- 2. RERANKER PERSONNALISÉ (Correction du Warning) ---
class MixedbreadReranker(BaseRetriever):
    retriever: Any 
    k: int = 4

    def _get_relevant_documents(self, query: str, *, run_manager: CallbackManagerForRetrieverRun) -> List[Document]:
        # 1. On récupère les documents dans ChromaDB (Local, donc toujours dispo !)
        initial_docs = self.retriever.invoke(query)
        
        try:
            # 2. Tentative de Reranking avec Mixedbread
            res = mxbai_client.reranking(
                model=MODEL_RRK,
                query=query,
                input=[doc.page_content for doc in initial_docs],
                top_k=self.k
            )
            
            final_docs = []
            for r in res.data:
                final_docs.append(Document(
                    page_content=r.input, 
                    metadata=initial_docs[r.index].metadata
                ))
            return final_docs
            
        except Exception as e:
            # 3. LE SECOURS : Si l'API 503 répond, on renvoie les docs bruts de Chroma
            print(f"⚠️ Mixedbread indisponible ({e}). Passage en mode secours.")
            # On renvoie juste les 'k' premiers documents de ChromaDB sans les retrier
            return initial_docs[:self.k]

# --- 3. INITIALISATION DU RAG TOOL ---
def init_rag_tool():
    db = get_vector_db()
    base_retriever = db.as_retriever(search_kwargs={"k": 25})
    reranker = MixedbreadReranker(retriever=base_retriever, k=5)
    
    llm = ChatOpenAI(model=LLM_NAME, temperature=0)
    
    system_prompt = (
        "Tu es un expert en actuariat. Utilise le contexte suivant pour répondre : \n\n {context}\n\n"
        "Si la réponse n'est pas dans le contexte, utilise tes connaissances générales ou Web_search."
    )
    prompt = ChatPromptTemplate.from_messages([
        ("system", system_prompt), 
        ("human", "{input}")
    ])
    
    qa_chain = create_stuff_documents_chain(llm, prompt)
    rag_chain = create_retrieval_chain(reranker, qa_chain)
    
    # Transformation de la chain en outil pour l'agent
    return rag_chain.as_tool(
        name="RAG_search",
        description="Cherche des informations techniques dans les mémoires d'actuariat."
    )

def init_websearch_tool():
    return DuckDuckGoSearchRun(
        name="Web_search",
        max_results=3,
        description="Recherche d'informations actuelles sur le web",
    )

def create_agent():
    rag_tool = init_rag_tool()
    web_search_tool = init_websearch_tool()
    memory = MemorySaver()  
    
    llm_4o = ChatOpenAI(model=LLM_NAME, temperature=0, streaming=True)
    tools = [rag_tool, web_search_tool]
    
    system_message = """
    Tu es l'Expert Actuariel Senior de Dataltist (Groupe Forsides). Ton rôle est de fournir des analyses de haute précision aux actuaires et professionnels de l'assurance.

    --- DIRECTIVES DE RECHERCHE ---
    1. PRIORITÉ LOCALE : Utilise toujours 'RAG_search' en premier pour les concepts théoriques, thèses, modèles mathématiques et mémoires actuariels. C'est ta base de connaissance de référence (800+ documents).
    2. ACTUALITÉ : Utilise 'Web_search' uniquement pour les réglementations parues après 2024, les actualités de l'ACPR/EIOPA de 2025/2026, ou les chiffres de marché récents.
    3. HYBRIDATION : Si une question lie théorie et actualité, utilise les deux outils pour une réponse complète.

    --- STRUCTURE DE RÉPONSE ---
    - Ta réponse doit être pédagogique, structurée avec des titres (##) et des points clés.
    - Ne fais jamais de réponses en un seul bloc de texte.
    - Utilise des tableaux Markdown si tu dois comparer des données.

    --- CITATIONS & SOURCES ---
    - SOURCE RAG : Cite impérativement le nom du fichier ou le titre du mémoire trouvé dans les métadonnées (ex: Source : Mémoire "Modélisation du risque cyber", 2024).
    - SOURCE WEB : Pour chaque info issue d'internet, affiche le nom du site et le lien entre parenthèses (ex: Source : [ACPR](https://acpr.banque-france.fr)).

    --- FORMATAGE MATHÉMATIQUE ---
    - STRICT : N'utilise JAMAIS de notation LaTeX (pas de $, pas de \, pas de [ ]).
    - Écris les formules de manière textuelle et lisible (ex: "Provision = Somme des Sinistres * Facteur de développement").
    - Utilise l'étoile (*) pour la multiplication et le chapeau (^) pour les puissances.

    --- TON & POSTURE ---
    - Sois direct, professionnel et pragmatique.
    - Si la base de données ne contient pas l'information, fait une recherche web ou une analyse basée sur les principes actuariels généraux.
    """

    return create_react_agent(
        llm_4o, 
        tools, 
        prompt=system_message, 
        checkpointer=memory
    )