ragflow_preprocess/agents/translation.py

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

"""
Agent LLM pour la traduction de contenu
"""

import json
import requests
from typing import Optional, Dict

from .base import LLMBaseAgent

class TranslationAgent(LLMBaseAgent):
    """
    Agent LLM spécialisé dans la traduction de texte
    """
    
    def __init__(self, model_name: str = "mistral", **config):
        """
        Initialise l'agent de traduction
        
        Args:
            model_name (str): Nom du modèle de traduction (défaut: mistral)
            **config: Paramètres de configuration supplémentaires
        """
        super().__init__(model_name, **config)
        
        # Définir les prompts de traduction
        self.prompts = {
            "fr_to_en": "Traduis le texte suivant du français vers l'anglais. "
                        "Préserve le formatage, le ton et le style du texte original. "
                        "Assure-toi que la traduction est fluide et naturelle.\n\n"
                        "Texte français :\n{text}",
            
            "en_to_fr": "Traduis le texte suivant de l'anglais vers le français. "
                        "Préserve le formatage, le ton et le style du texte original. "
                        "Assure-toi que la traduction est fluide et naturelle.\n\n"
                        "Texte anglais :\n{text}"
        }
    
    def generate(self, text: str, source_lang: str = "fr", target_lang: str = "en") -> str:
        """
        Traduit un texte d'une langue source vers une langue cible
        
        Args:
            text (str): Texte à traduire
            source_lang (str): Langue source (fr ou en)
            target_lang (str): Langue cible (fr ou en)
            
        Returns:
            str: La traduction générée par le modèle
        """
        if not text:
            return ""
        
        # Déterminer la direction de traduction
        if source_lang == "fr" and target_lang == "en":
            direction = "fr_to_en"
        elif source_lang == "en" and target_lang == "fr":
            direction = "en_to_fr"
        else:
            return f"Traduction non prise en charge: {source_lang} vers {target_lang}"
        
        # Construire le prompt
        prompt = self.prompts[direction].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.1),  # Basse pour traduction précise
                    "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 une traduction simulée
            if direction == "fr_to_en":
                if "schéma" in text.lower():
                    return text.replace("schéma", "diagram").replace("Schéma", "Diagram")
                elif "tableau" in text.lower():
                    return text.replace("tableau", "table").replace("Tableau", "Table")
                elif "formule" in text.lower():
                    return text.replace("formule", "formula").replace("Formule", "Formula")
                else:
                    # Exemple très simplifié de "traduction"
                    translations = {
                        "Le": "The", "la": "the", "les": "the", "des": "the",
                        "et": "and", "ou": "or", "pour": "for", "avec": "with",
                        "est": "is", "sont": "are", "contient": "contains",
                        "montre": "shows", "représente": "represents",
                        "plusieurs": "several", "important": "important",
                        "information": "information", "données": "data",
                        "processus": "process", "système": "system",
                        "analyse": "analysis", "résultat": "result"
                    }
                    
                    # Remplacement simple mot à mot
                    result = text
                    for fr, en in translations.items():
                        result = result.replace(f" {fr} ", f" {en} ")
                    
                    return result
            else:  # en_to_fr
                if "diagram" in text.lower():
                    return text.replace("diagram", "schéma").replace("Diagram", "Schéma")
                elif "table" in text.lower():
                    return text.replace("table", "tableau").replace("Table", "Tableau")
                elif "formula" in text.lower():
                    return text.replace("formula", "formule").replace("Formula", "Formule")
                else:
                    # Exemple très simplifié de "traduction"
                    translations = {
                        "The": "Le", "the": "le", "and": "et", "or": "ou",
                        "for": "pour", "with": "avec", "is": "est",
                        "are": "sont", "contains": "contient", "shows": "montre",
                        "represents": "représente", "several": "plusieurs",
                        "important": "important", "information": "information",
                        "data": "données", "process": "processus",
                        "system": "système", "analysis": "analyse",
                        "result": "résultat"
                    }
                    
                    # Remplacement simple mot à mot
                    result = text
                    for en, fr in translations.items():
                        result = result.replace(f" {en} ", f" {fr} ")
                    
                    return result
        
        except Exception as e:
            return f"Erreur lors de la traduction: {str(e)}"