ragflow_preprocess/agents/translation.py

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

"""
Agent pour la traduction 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 TranslationAgent(LLMBaseAgent):
    """
    Agent pour la traduction de texte
    """
    
    def __init__(self, model_name: str, endpoint: str = "http://217.182.105.173:11434", **config):
        """
        Initialise l'agent de traduction
        
        Args:
            model_name (str): Nom du modèle à utiliser
            endpoint (str): URL de l'API Ollama
            **config: Configuration supplémentaire
        """
        super().__init__(model_name, endpoint, **config)
        
        # Configuration par défaut pour la traduction
        default_config = {
            "source_language": "en",
            "target_language": "fr"
        }
        
        # 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 traduction
        self.log_dir = os.path.join("data", "translations")
        os.makedirs(self.log_dir, exist_ok=True)
    
    def generate(self, prompt: str, images: Optional[List[bytes]] = None, 
                 source_language: Optional[str] = None, target_language: Optional[str] = None) -> str:
        """
        Traduit un texte
        
        Args:
            prompt (str): Texte à traduire
            images (List[bytes], optional): Non utilisé pour la traduction
            source_language (str, optional): Langue source (par défaut: celle de la configuration)
            target_language (str, optional): Langue cible (par défaut: celle de la configuration)
            
        Returns:
            str: Traduction générée
        """
        if not prompt or not prompt.strip():
            return "Erreur: Aucun texte fourni pour la traduction"
        
        # Utiliser les langues spécifiées ou celles de la configuration
        src_lang = source_language or self.config["source_language"]
        tgt_lang = target_language or self.config["target_language"]
        
        # Génération d'un ID unique pour cette traduction
        translation_id = str(uuid.uuid4())[:8]
        timestamp = time.strftime("%Y%m%d-%H%M%S")
        
        # Construire le prompt pour le modèle
        system_prompt = f"You are a professional and accurate translator. Translate the following text "
        system_prompt += f"from {src_lang} to {tgt_lang}. Maintain the original formatting and structure as much as possible."
        
        # Construire le message utilisateur avec le texte à traduire
        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 traduction 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.1),  # Température plus basse pour la traduction
                        "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.1),
                        "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 traduction dans un fichier
            translation_path = os.path.join(self.log_dir, f"{timestamp}_{translation_id}.txt")
            with open(translation_path, "w", encoding="utf-8") as f:
                f.write(f"Source language: {src_lang}\n")
                f.write(f"Target language: {tgt_lang}\n")
                f.write(f"Model: {self.model_name}\n\n")
                f.write(f"Original text:\n{prompt}\n\n")
                f.write(f"Translation:\n{result}")
            
            print(f"Traduction enregistrée dans: {translation_path}")
            return result
            
        except Exception as e:
            error_msg = f"Erreur lors de la traduction: {str(e)}"
            print(error_msg)
            
            # Enregistrer l'erreur
            error_path = os.path.join(self.log_dir, f"{timestamp}_{translation_id}_error.txt")
            with open(error_path, "w", encoding="utf-8") as f:
                f.write(f"Source language: {src_lang}\n")
                f.write(f"Target language: {tgt_lang}\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