ragflow_preprocess/agents/summary.py

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

"""
Agent pour la synthèse et le résumé de texte
"""

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

from .base import LLMBaseAgent
from utils.api_ollama import OllamaAPI

class SummaryAgent(LLMBaseAgent):
    """
    Agent pour la synthèse et le résumé de texte
    """
    
    def __init__(self, model_name: str, endpoint: str = "http://217.182.105.173:11434", **config):
        """
        Initialise l'agent de résumé
        
        Args:
            model_name (str): Nom du modèle à utiliser
            endpoint (str): URL de l'API Ollama
            **config: Configuration supplémentaire
        """
        # Appeler le constructeur de la classe parent avec les paramètres requis
        super().__init__(model_name, endpoint, **config)
        
        # Configuration par défaut pour les résumés
        default_config = {
            "language": "fr",
            "summary_length": "medium"  # 'short', 'medium', 'long'
        }
        
        # Mettre à jour la configuration avec les valeurs par défaut si elles ne sont pas spécifiées
        for key, value in default_config.items():
            if key not in self.config:
                self.config[key] = value
        
        # Création du répertoire pour les journaux de résumé
        self.log_dir = os.path.join("data", "summaries")
        os.makedirs(self.log_dir, exist_ok=True)
    
    def generate(self, prompt: str, images: Optional[List[bytes]] = None, 
                 summary_length: Optional[str] = None, language: Optional[str] = None) -> str:
        """
        Résume un texte
        
        Args:
            prompt (str): Texte à résumer
            images (List[bytes], optional): Non utilisé pour le résumé
            summary_length (str, optional): Longueur du résumé ('short', 'medium', 'long')
            language (str, optional): Langue du résumé
            
        Returns:
            str: Résumé généré
        """
        if not prompt or not prompt.strip():
            return "Erreur: Aucun texte fourni pour le résumé"
        
        # Utiliser les paramètres spécifiés ou ceux de la configuration
        length = summary_length or self.config["summary_length"]
        lang = language or self.config["language"]
        
        # Génération d'un ID unique pour ce résumé
        summary_id = str(uuid.uuid4())[:8]
        timestamp = time.strftime("%Y%m%d-%H%M%S")
        
        # Définir les instructions de longueur en fonction de la langue
        length_instructions = {
            "short": {
                "fr": "Fais un résumé court et concis (environ 2-3 phrases).",
                "en": "Create a short and concise summary (about 2-3 sentences)."
            },
            "medium": {
                "fr": "Fais un résumé de taille moyenne (environ 1-2 paragraphes).",
                "en": "Create a medium-length summary (about 1-2 paragraphs)."
            },
            "long": {
                "fr": "Fais un résumé détaillé (environ 3-4 paragraphes).",
                "en": "Create a detailed summary (about 3-4 paragraphs)."
            }
        }
        
        # Obtenir l'instruction de longueur dans la langue appropriée
        length_instruction = length_instructions.get(length, length_instructions["medium"]).get(lang, length_instructions["medium"]["fr"])
        
        # Construire le prompt pour le modèle
        if lang == "fr":
            system_prompt = f"Tu es un expert en synthèse et résumé de textes. {length_instruction} "
            system_prompt += "Conserve l'information essentielle et la structure logique du texte original."
        else:
            system_prompt = f"You are an expert in summarizing texts. {length_instruction} "
            system_prompt += "Preserve the essential information and logical structure of the original text."
        
        # Construire le message utilisateur avec le texte à résumer
        user_prompt = prompt
        
        # Créer l'API Ollama pour l'appel direct
        api = OllamaAPI(base_url=self.endpoint)
        
        # Journaliser le prompt complet
        full_prompt = f"System: {system_prompt}\n\nUser: {user_prompt}"
        print(f"Envoi du prompt de résumé au modèle {self.model_name}:\n{system_prompt}")
        
        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
                messages = [
                    {"role": "system", "content": system_prompt},
                    {"role": "user", "content": user_prompt}
                ]
                
                response = api.chat(
                    model=self.model_name,
                    messages=messages,
                    options={
                        "temperature": self.config.get("temperature", 0.2),
                        "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
                prompt_text = f"{system_prompt}\n\n{user_prompt}"
                
                response = api.generate(
                    model=self.model_name,
                    prompt=prompt_text,
                    options={
                        "temperature": self.config.get("temperature", 0.2),
                        "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 le résumé dans un fichier
            summary_path = os.path.join(self.log_dir, f"{timestamp}_{summary_id}.txt")
            with open(summary_path, "w", encoding="utf-8") as f:
                f.write(f"Language: {lang}\n")
                f.write(f"Summary length: {length}\n")
                f.write(f"Model: {self.model_name}\n\n")
                f.write(f"Original text:\n{prompt}\n\n")
                f.write(f"Summary:\n{result}")
            
            print(f"Résumé enregistré dans: {summary_path}")
            return result
            
        except Exception as e:
            error_msg = f"Erreur lors de la génération du résumé: {str(e)}"
            print(error_msg)
            
            # Enregistrer l'erreur
            error_path = os.path.join(self.log_dir, f"{timestamp}_{summary_id}_error.txt")
            with open(error_path, "w", encoding="utf-8") as f:
                f.write(f"Language: {lang}\n")
                f.write(f"Summary length: {length}\n")
                f.write(f"Model: {self.model_name}\n\n")
                f.write(f"Original text:\n{prompt}\n\n")
                f.write(f"Error:\n{str(e)}")
                
            return error_msg