ragflow_preprocess/agents/rewriter.py

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

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

import requests
from typing import Dict, Optional, Union

from .base import LLMBaseAgent

class RewriterAgent(LLMBaseAgent):
    """
    Agent LLM spécialisé dans la reformulation et l'adaptation de texte
    """
    
    def __init__(self, model_name: str = "mistral", **config):
        """
        Initialise l'agent de reformulation
        
        Args:
            model_name (str): Nom du modèle de reformulation (défaut: mistral)
            **config: Paramètres de configuration supplémentaires
        """
        super().__init__(model_name, **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}"
            }
        }
    
    def generate(self, text: str, mode: str = "rag", custom_prompt: Optional[str] = "") -> str:
        """
        Reformule un texte selon le mode spécifié
        
        Args:
            text (str): Texte à reformuler
            mode (str): Mode de reformulation (simplifier, détailler, rag, formal, bullet)
            custom_prompt (str, optional): Prompt personnalisé pour la reformulation
            
        Returns:
            str: Le texte reformulé
        """
        if not text:
            return ""
        
        # Déterminer la langue et le prompt à utiliser
        lang = self.config.get("language", "fr")
        
        if custom_prompt:
            prompt = custom_prompt.format(text=text)
        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]
            prompt = prompt_template.format(text=text)
        
        try:
            # Construire la payload pour l'API Ollama
            payload = {
                "model": self.model_name,
                "prompt": prompt,
                "options": {
                    "temperature": self.config.get("temperature", 0.3),  # Légèrement plus créatif pour la reformulation
                    "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)
                }
            }
            
            # Dans une implémentation réelle, envoyer la requête à l'API Ollama
            # response = requests.post(f"{self.endpoint}/api/generate", json=payload)
            # json_response = response.json()
            # return json_response.get("response", "")
            
            # Pour cette démonstration, retourner des exemples de reformulation
            if mode == "simplifier" or mode == "simplify":
                return ("Ce texte a été simplifié pour être plus facile à comprendre. "
                        "Les mots compliqués ont été remplacés par des mots plus simples. "
                        "Les phrases longues ont été raccourcies. Les idées principales "
                        "restent les mêmes, mais elles sont expliquées plus clairement.")
            
            elif mode == "détailler" or mode == "elaborate":
                return (f"{text}\n\nEn outre, il est important de noter que ce contenu s'inscrit "
                        "dans un contexte plus large. Plusieurs exemples concrets illustrent ce point : "
                        "premièrement, l'application pratique de ces concepts dans des situations réelles; "
                        "deuxièmement, les différentes interprétations possibles selon le domaine d'expertise; "
                        "et troisièmement, les implications à long terme de ces informations. "
                        "Cette perspective élargie permet une compréhension plus approfondie du sujet.")
            
            elif mode == "formal":
                return ("Il convient de préciser que le contenu susmentionné présente des caractéristiques "
                        "particulièrement pertinentes dans le cadre de l'analyse proposée. En effet, "
                        "l'examen minutieux des éléments constitutifs révèle une structure cohérente "
                        "dont la logique sous-jacente manifeste une organisation méthodique des concepts. "
                        "Par conséquent, il est possible d'affirmer que les principes énoncés "
                        "s'inscrivent dans un paradigme rigoureux qui mérite une attention scientifique.")
            
            elif mode == "bullet":
                # Transformation en liste à puces
                bullet_points = ["• Point principal: Le contenu présente des informations essentielles",
                                 "• Structure: Organisation logique des éléments clés",
                                 "• Application: Utilisations pratiques dans divers contextes",
                                 "• Avantages: Amélioration de la compréhension et de l'efficacité",
                                 "• Limitations: Considérations importantes à prendre en compte"]
                
                return "\n".join(bullet_points)
            
            else:  # rag par défaut
                return ("Ce contenu optimisé pour les systèmes RAG contient des mots-clés pertinents "
                        "et une structure sémantique améliorée. Les concepts principaux sont clairement "
                        "définis et leurs relations sont explicitement établies. Les informations "
                        "sont présentées de manière à faciliter la recherche et la récupération "
                        "automatisées, avec une organisation logique qui met en évidence les "
                        "éléments essentiels du sujet traité. Chaque section est conçue pour "
                        "maximiser la pertinence lors des requêtes d'information.")
        
        except Exception as e:
            return f"Erreur lors de la reformulation: {str(e)}"