ragflow_preprocess/utils/markdown_export.py

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

"""
Module pour l'exportation des résultats au format Markdown
"""

import os
import base64
import re
from typing import List, Dict, Any, Optional, Union


class MarkdownExporter:
    """
    Classe pour exporter les résultats d'analyse en Markdown
    """
    
    def __init__(self, output_dir: Optional[str] = "") -> None:
        """
        Initialise l'exporteur Markdown
        
        Args:
            output_dir (str, optional): Répertoire de sortie pour les fichiers générés
        """
        if not output_dir:
            self.output_dir = os.path.join(os.getcwd(), "data", "outputs")
        else:
            self.output_dir = output_dir
            
        os.makedirs(self.output_dir, exist_ok=True)
        
        # Compteur pour les images exportées
        self.image_counter = 0
    
    def create_markdown_from_selections(self, selections: List[Dict[str, Any]], 
                                      analysis_results: Dict[int, str],
                                      document_title: str = "Analyse de document",
                                      include_images: bool = True) -> str:
        """
        Crée un document Markdown à partir des sélections et des résultats d'analyse
        
        Args:
            selections (List[Dict]): Liste des sélections (régions d'intérêt)
            analysis_results (Dict): Dictionnaire des résultats d'analyse par ID de sélection
            document_title (str): Titre du document Markdown
            include_images (bool): Si True, inclut les images dans le Markdown
            
        Returns:
            str: Contenu du document Markdown généré
        """
        # Trier les sélections par numéro de page
        sorted_selections = sorted(selections, key=lambda x: x.get("page", 0))
        
        # Commencer le document avec le titre
        markdown_content = f"# {document_title}\n\n"
        
        # Ajouter chaque sélection
        for selection in sorted_selections:
            page = selection.get("page", 0) + 1
            selection_type = selection.get("type", "zone")
            context = selection.get("context", "")
            
            # Titre de la section
            markdown_content += f"## {selection_type.capitalize()} - Page {page}\n\n"
            
            # Ajouter l'image si demandé et disponible
            if include_images and "image_data" in selection:
                image_path = self._save_image(selection["image_data"])
                if image_path:
                    rel_path = os.path.relpath(image_path, self.output_dir)
                    markdown_content += f"![{selection_type} de la page {page}]({rel_path})\n\n"
            
            # Ajouter le contexte
            if context:
                markdown_content += f"**Contexte** :\n{context}\n\n"
            
            # Ajouter les résultats d'analyse
            selection_id = id(selection)
            if selection_id in analysis_results:
                markdown_content += f"**Analyse IA** :\n{analysis_results[selection_id]}\n\n"
            
            # Ajouter un séparateur
            markdown_content += "---\n\n"
        
        return markdown_content
    
    def export_to_file(self, markdown_content: str, filename: str = "analyse_document.md") -> str:
        """
        Exporte le contenu Markdown dans un fichier
        
        Args:
            markdown_content (str): Contenu Markdown à exporter
            filename (str): Nom du fichier de sortie
            
        Returns:
            str: Chemin absolu du fichier créé
        """
        file_path = os.path.join(self.output_dir, filename)
        
        try:
            with open(file_path, "w", encoding="utf-8") as f:
                f.write(markdown_content)
            
            return file_path
        
        except Exception as e:
            print(f"Erreur lors de l'exportation du fichier Markdown: {str(e)}")
            return ""
    
    def _save_image(self, image_data: bytes, prefix: str = "image") -> Optional[str]:
        """
        Sauvegarde une image et retourne son chemin
        
        Args:
            image_data (bytes): Données de l'image
            prefix (str): Préfixe pour le nom du fichier
            
        Returns:
            Optional[str]: Chemin de l'image sauvegardée ou None en cas d'erreur
        """
        try:
            # Incrémenter le compteur d'images
            self.image_counter += 1
            
            # Créer le répertoire d'images s'il n'existe pas
            images_dir = os.path.join(self.output_dir, "images")
            os.makedirs(images_dir, exist_ok=True)
            
            # Chemin de l'image
            image_path = os.path.join(images_dir, f"{prefix}_{self.image_counter}.png")
            
            # Sauvegarder l'image
            with open(image_path, "wb") as f:
                f.write(image_data)
            
            return image_path
        
        except Exception as e:
            print(f"Erreur lors de la sauvegarde de l'image: {str(e)}")
            return None
    
    def format_analysis_result(self, result: str, include_parameters: bool = True) -> str:
        """
        Formate un résultat d'analyse pour le Markdown
        
        Args:
            result (str): Texte du résultat d'analyse
            include_parameters (bool): Si True, inclut la section des paramètres
            
        Returns:
            str: Résultat formaté pour le Markdown
        """
        # Simple nettoyage des sections
        formatted = result
        
        # Séparer les paramètres
        if not include_parameters and "**Paramètres utilisés**" in result:
            formatted = result.split("**Paramètres utilisés**")[0].strip()
        
        return formatted

    @staticmethod
    def sanitize_filename(title: str) -> str:
        """
        Convertit un titre en nom de fichier sécurisé
        
        Args:
            title (str): Titre à convertir
            
        Returns:
            str: Nom de fichier sécurisé
        """
        # Remplacer les caractères non alphanumériques par des tirets
        sanitized = re.sub(r'[^a-zA-Z0-9]', '-', title.lower())
        # Réduire les tirets multiples à un seul
        sanitized = re.sub(r'-+', '-', sanitized)
        # Limiter la longueur et supprimer les tirets au début/fin
        return sanitized[:50].strip('-')
    
    def embed_images_in_markdown(self, markdown_content: str) -> str:
        """
        Remplace les références d'images par des images en base64 intégrées
        
        Args:
            markdown_content (str): Contenu Markdown avec références d'images
            
        Returns:
            str: Contenu Markdown avec images intégrées
        """
        # Regex pour trouver les références d'images
        img_pattern = r'!\[(.*?)\]\((.*?)\)'
        
        def replace_image(match):
            alt_text = match.group(1)
            img_path = match.group(2)
            
            # Résoudre le chemin relatif
            if not os.path.isabs(img_path):
                img_path = os.path.join(self.output_dir, img_path)
            
            try:
                # Lire l'image et l'encoder en base64
                with open(img_path, "rb") as img_file:
                    img_data = img_file.read()
                    img_base64 = base64.b64encode(img_data).decode('utf-8')
                
                # Déterminer le type MIME
                if img_path.lower().endswith('.png'):
                    mime_type = 'image/png'
                elif img_path.lower().endswith(('.jpg', '.jpeg')):
                    mime_type = 'image/jpeg'
                else:
                    mime_type = 'image/png'  # par défaut
                
                # Retourner l'image intégrée en base64
                return f'![{alt_text}](data:{mime_type};base64,{img_base64})'
            
            except Exception as e:
                print(f"Erreur lors de l'intégration de l'image {img_path}: {str(e)}")
                return match.group(0)  # Conserver l'original en cas d'erreur
        
        # Remplacer toutes les références d'images
        return re.sub(img_pattern, replace_image, markdown_content)