ragflow_preprocess/agents/rewriter.py

#!/usr/bin/env python3
# -*- coding: utf-8 -*-

"""
Agent LLM pour la reformulation et l'adaptation de contenu
"""

import os
import time
import uuid
from typing import Dict, Optional, Union, List, Any

from .base import LLMBaseAgent
from utils.api_ollama import OllamaAPI

class RewriterAgent(LLMBaseAgent):
    """
    Agent LLM spécialisé dans la reformulation et l'adaptation de texte
    """
    
    def __init__(self, model_name: str, endpoint: str = "http://217.182.105.173:11434", **config):
        """
        Initialise l'agent de reformulation
        
        Args:
            model_name (str): Nom du modèle de reformulation
            endpoint (str): URL de l'API Ollama
            **config: Paramètres de configuration supplémentaires
        """
        super().__init__(model_name, endpoint, **config)
        
        # Définir les modes de reformulation et leurs prompts
        self.modes = {
            "fr": {
                "simplifier": "Reformule le texte suivant pour le rendre plus accessible "
                              "et facile à comprendre. Simplifie le vocabulaire et la structure "
                              "des phrases tout en préservant le sens original.\n\n"
                              "Texte original :\n{text}",
                
                "détailler": "Développe et enrichis le texte suivant en ajoutant des détails, "
                             "des explications et des exemples pertinents. Garde le même ton "
                             "et le même style, mais rends le contenu plus complet.\n\n"
                             "Texte à développer :\n{text}",
                
                "rag": "Reformule le texte suivant pour l'optimiser pour un système de RAG "
                       "(Retrieval Augmented Generation). Assure-toi qu'il contient des mots-clés "
                       "pertinents, qu'il est bien structuré pour la recherche sémantique, et "
                       "qu'il présente clairement les informations essentielles.\n\n"
                       "Texte à optimiser :\n{text}",
                
                "formal": "Reformule le texte suivant dans un style plus formel et académique. "
                          "Utilise un vocabulaire précis, une structure de phrase soignée, et "
                          "un ton professionnel tout en préservant le contenu original.\n\n"
                          "Texte à formaliser :\n{text}",
                
                "bullet": "Transforme le texte suivant en une liste à puces claire et concise. "
                          "Extrais les points essentiels et présente-les dans un format facile "
                          "à lire et à comprendre.\n\n"
                          "Texte à transformer :\n{text}"
            },
            "en": {
                "simplify": "Rewrite the following text to make it more accessible "
                            "and easy to understand. Simplify vocabulary and sentence structure "
                            "while preserving the original meaning.\n\n"
                            "Original text:\n{text}",
                
                "elaborate": "Expand and enrich the following text by adding relevant details, "
                             "explanations, and examples. Keep the same tone and style, "
                             "but make the content more comprehensive.\n\n"
                             "Text to expand:\n{text}",
                
                "rag": "Rewrite the following text to optimize it for a RAG "
                       "(Retrieval Augmented Generation) system. Ensure it contains relevant "
                       "keywords, is well-structured for semantic search, and "
                       "clearly presents essential information.\n\n"
                       "Text to optimize:\n{text}",
                
                "formal": "Rewrite the following text in a more formal and academic style. "
                          "Use precise vocabulary, careful sentence structure, and "
                          "a professional tone while preserving the original content.\n\n"
                          "Text to formalize:\n{text}",
                
                "bullet": "Transform the following text into a clear and concise bullet point list. "
                          "Extract the essential points and present them in a format that is easy "
                          "to read and understand.\n\n"
                          "Text to transform:\n{text}"
            }
        }
        
        # Création du répertoire pour les journaux de reformulation
        self.log_dir = os.path.join("data", "rewrites")
        os.makedirs(self.log_dir, exist_ok=True)
    
    def generate(self, text: str, images: Optional[List[bytes]] = None, mode: str = "rag", 
                 custom_prompt: Optional[str] = "", language: Optional[str] = None) -> str:
        """
        Reformule un texte selon le mode spécifié
        
        Args:
            text (str): Texte à reformuler
            images (List[bytes], optional): Non utilisé pour la reformulation
            mode (str): Mode de reformulation (simplifier, détailler, rag, formal, bullet)
            custom_prompt (str, optional): Prompt personnalisé pour la reformulation
            language (str, optional): Langue du texte à reformuler
            
        Returns:
            str: Le texte reformulé
        """
        if not text or not text.strip():
            return "Erreur: Aucun texte fourni pour la reformulation"
        
        # Déterminer la langue et le prompt à utiliser
        lang = language or self.config.get("language", "fr")
        
        # Génération d'un ID unique pour cette reformulation
        rewrite_id = str(uuid.uuid4())[:8]
        timestamp = time.strftime("%Y%m%d-%H%M%S")
        
        if custom_prompt:
            prompt_template = custom_prompt
            mode_name = "custom"
        else:
            # Vérifier que le mode existe pour la langue spécifiée
            if lang not in self.modes or mode not in self.modes[lang]:
                # Fallback sur français et mode RAG si non disponible
                lang = "fr"
                mode = "rag"
            
            prompt_template = self.modes[lang][mode]
            mode_name = mode
        
        # Formater le prompt avec le texte
        prompt = prompt_template.format(text=text)
        
        # Créer l'API Ollama pour l'appel direct
        api = OllamaAPI(base_url=self.endpoint)
        
        # Journaliser le prompt complet
        print(f"Envoi du prompt de reformulation (mode {mode_name}) au modèle {self.model_name}")
        
        try:
            # Pour les modèles qui supportent le format de chat
            if any(name in self.model_name.lower() for name in ["llama", "mistral", "deepseek", "qwen"]):
                # Formater en tant que messages de chat
                system_prompt = f"Tu es un expert en reformulation de texte spécialisé dans le mode '{mode_name}'."
                
                messages = [
                    {"role": "system", "content": system_prompt},
                    {"role": "user", "content": prompt}
                ]
                
                response = api.chat(
                    model=self.model_name,
                    messages=messages,
                    options={
                        "temperature": self.config.get("temperature", 0.3),
                        "top_p": self.config.get("top_p", 0.95),
                        "top_k": self.config.get("top_k", 40),
                        "num_predict": self.config.get("max_tokens", 2048)
                    }
                )
                
                if "message" in response and "content" in response["message"]:
                    result = response["message"]["content"]
                else:
                    result = response.get("response", "Erreur: Format de réponse inattendu")
            else:
                # Format de génération standard pour les autres modèles
                response = api.generate(
                    model=self.model_name,
                    prompt=prompt,
                    options={
                        "temperature": self.config.get("temperature", 0.3),
                        "top_p": self.config.get("top_p", 0.95),
                        "top_k": self.config.get("top_k", 40),
                        "num_predict": self.config.get("max_tokens", 2048)
                    }
                )
                
                result = response.get("response", "Erreur: Pas de réponse")
            
            # Enregistrer la reformulation dans un fichier
            rewrite_path = os.path.join(self.log_dir, f"{timestamp}_{rewrite_id}.txt")
            with open(rewrite_path, "w", encoding="utf-8") as f:
                f.write(f"Language: {lang}\n")
                f.write(f"Mode: {mode_name}\n")
                f.write(f"Model: {self.model_name}\n\n")
                f.write(f"Original text:\n{text}\n\n")
                f.write(f"Rewritten text:\n{result}")
            
            print(f"Reformulation enregistrée dans: {rewrite_path}")
            return result
            
        except Exception as e:
            error_msg = f"Erreur lors de la reformulation: {str(e)}"
            print(error_msg)
            
            # Enregistrer l'erreur
            error_path = os.path.join(self.log_dir, f"{timestamp}_{rewrite_id}_error.txt")
            with open(error_path, "w", encoding="utf-8") as f:
                f.write(f"Language: {lang}\n")
                f.write(f"Mode: {mode_name}\n")
                f.write(f"Model: {self.model_name}\n\n")
                f.write(f"Original text:\n{text}\n\n")
                f.write(f"Error:\n{str(e)}")
                
            return error_msg