AIagent/llm_classes/rag.py

import json
import datetime
import requests
from typing import Dict, List, Any, Optional, Union, Tuple

class RAG:
    """
    Classe pour l'intégration avec l'API RAG (Retrieval Augmented Generation)
    Cette classe ne dépend pas de la classe LLM et fonctionne de manière autonome
    """
    
    def __init__(self):
        """
        Initialisation des attributs de RAG
        """
        # Attributs publics constants (transformés en attributs d'instance)
        self.sessionID: str = ""
        
        # Attributs privés
        self._reqRAG: requests.Session = requests.Session()
        self._taContenu: Dict[str, Any] = {}
        self._m_sChatID: str = ""
        self._mn_Question: int = 0
        
        # Initialisation par défaut
        self._reqRAG.headers["Authorization"] = "Bearer ragflow-c4YTNkMzcwZDM1ODExZWZiODA2MDI0Mm"
        self._reqRAG.headers["Content-Type"] = "application/json"
    
    def _urlBase(self) -> str:
        """
        Retourne l'URL de base de l'API RAG
        """
        return "http://10.103.0.100/api/v1/chats/"
    
    def _chatID(self) -> str:
        """
        Retourne l'ID de chat actuel ou un ID par défaut
        """
        if self._m_sChatID:
            return self._m_sChatID
        else:
            return "ffb1058ed4b811ef8a900242ac120003"
    
    def _set_chatID(self, valeur: str) -> None:
        """
        Définit l'ID de chat
        """
        self._m_sChatID = valeur
    
    def _Contenu(self) -> str:
        """
        Sérialise le contenu en JSON
        """
        return json.dumps(self._taContenu)
    
    def _urlChat(self) -> str:
        """
        Retourne l'URL complète pour le chat
        """
        return f"{self._urlBase()}{self._chatID()}/"
    
    def __del__(self):
        """
        Destructeur - supprime la session de chat si nécessaire
        """
        # Équivalent de: SI mn_Question = 0 OU EnModeTest() ALORS supprimeSessionChat()
        if self._mn_Question == 0:
            self._supprimeSessionChat()
    
    def _creeSessionChat(self) -> Tuple[bool, str]:
        """
        Crée une nouvelle session de chat
        Retourne un tuple (succès, message)
        """
        self._taContenu.clear()
        
        # Préparation de la requête - sans utiliser l'attribut url directement
        req_url = f"{self._urlChat()}sessions"
        
        # Nom de la session basé sur la date et l'heure actuelles
        current_datetime = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
        self._taContenu["name"] = current_datetime
        
        try:
            response = self._reqRAG.post(
                url=req_url,
                json=self._taContenu
            )
            
            data = json.loads(response.text)
            
            if data.get("code") == 0:
                self.sessionID = data["data"]["id"]
                return True, self._formateReponse(data["data"]["messages"][0]["content"])
            else:
                return False, self._formateReponse(data.get("message", "Erreur inconnue"))
                
        except Exception as e:
            return False, self._formateReponse(str(e))
    
    def Chat(self, question: str) -> str:
        """
        Envoie une question au chat RAG
        """
        # Crée une session si nécessaire
        if not self.sessionID:
            success, response = self._creeSessionChat()
            if not success:
                return self._formateReponse(response)
        
        # Incrémente le compteur de questions
        self._mn_Question += 1
        
        # Prépare la requête
        self._taContenu.clear()
        self._taContenu["question"] = question
        self._taContenu["stream"] = False
        self._taContenu["session_id"] = self.sessionID
        
        try:
            response = self._reqRAG.post(
                url=self._urlChat() + "completions",
                json=self._taContenu
            )
            
            data = json.loads(response.text)
            
            if data.get("code") == 0:
                return self._formateReponse(data["data"]["answer"])
            else:
                return self._formateReponse(data.get("message", "Erreur inconnue"))
                
        except Exception as e:
            return self._formateReponse(str(e))
    
    def _formateReponse(self, reponse: str) -> str:
        """
        Formate la réponse en HTML
        """
        # Conversion en HTML avec style personnalisé
        html = f"""<div style="background-color: #e8f4ff; border: 1px solid #c2e0ff; border-radius: 10px; padding: 10px 15px; font-family: Arial, sans-serif; font-size: 14px; color: #000; display: inline-block;">
        {reponse}
        </div>"""
        
        return html
    
    @staticmethod
    def formateQuestion(question: str) -> str:
        """
        Formate la question en HTML
        """
        # Conversion en HTML avec style personnalisé
        html = f"""<div style="text-align: right; font-family: Arial, sans-serif; font-size: 14px; color: #000; padding: 10px 15px;">
        {question}
        </div>"""
        
        return html
    
    def _supprimeSessionChat(self) -> str:
        """
        Supprime la session de chat actuelle
        """
        if not self.sessionID:
            return "Aucune session à supprimer"
            
        self._taContenu.clear()
        self._taContenu["ids"] = [self.sessionID]
        
        try:
            response = self._reqRAG.delete(
                url=self._urlChat() + "/sessions",
                json=self._taContenu
            )
            
            data = json.loads(response.text)
            
            if data.get("code") == 0:
                return "OK"
            else:
                return data.get("message", "Erreur inconnue")
                
        except Exception as e:
            return str(e)