llm_ticket3/test_orchestrator.py

from orchestrator import Orchestrator
from agents.agent_json_analyser import AgentJsonAnalyser
from agents.agent_image_sorter import AgentImageSorter
from agents.agent_image_analyser import AgentImageAnalyser
from agents.agent_report_generator import AgentReportGenerator

# Utilisation des modèles Medium pour les tests
from llm_classes.mistral_medium import MistralMedium
from llm_classes.pixtral_12b import Pixtral12b

import os
import logging
import sys
import time

# Configuration du logging
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s', 
                    filename='test_orchestrator.log', filemode='w')
logger = logging.getLogger("TestOrchestrator")

def test_orchestrator(ticket_id=None):
    """
    Exécute l'orchestrateur avec les agents définis
    
    Args:
        ticket_id: Identifiant du ticket à traiter (optionnel)
    """
    # Vérifier que le dossier output existe
    if not os.path.exists("output/"):
        os.makedirs("output/")
        logger.warning("Le dossier output/ n'existait pas et a été créé")
        print("ATTENTION: Le dossier output/ n'existait pas et a été créé")
    
    # Vérifier le contenu du dossier output
    tickets = [d for d in os.listdir("output/") if d.startswith("ticket_") and os.path.isdir(os.path.join("output/", d))]
    logger.info(f"Tickets trouvés dans output/: {len(tickets)}")
    print(f"Tickets existants dans output/: {len(tickets)}")
    
    if len(tickets) == 0:
        logger.error("Aucun ticket trouvé dans le dossier output/")
        print("ERREUR: Aucun ticket trouvé dans le dossier output/")
        return
    
    # Initialisation des LLM
    print("Initialisation des modèles LLM...")
    
    start_time = time.time()
    
    # Utilisation de Mistral Medium pour l'analyse JSON et la génération de rapports
    json_llm = MistralMedium()
    logger.info("LLM MistralMedium initialisé pour l'analyse JSON")
    
    # Utilisation de Pixtral12b pour le tri et l'analyse d'images
    image_sorter_llm = Pixtral12b()
    logger.info("LLM Pixtral12b initialisé pour le tri d'images")
    
    image_analyser_llm = Pixtral12b()
    logger.info("LLM Pixtral12b initialisé pour l'analyse d'images")
    
    report_generator_llm = MistralMedium()
    logger.info("LLM MistralMedium initialisé pour la génération de rapports")
    
    llm_init_time = time.time() - start_time
    print(f"Tous les modèles LLM ont été initialisés en {llm_init_time:.2f} secondes")
    
    # Création des agents
    print("Création des agents...")
    json_agent = AgentJsonAnalyser(json_llm)
    image_sorter = AgentImageSorter(image_sorter_llm)
    image_analyser = AgentImageAnalyser(image_analyser_llm)
    report_generator = AgentReportGenerator(report_generator_llm)
    
    print("Tous les agents ont été créés")
    
    # Initialisation de l'orchestrateur avec les agents
    logger.info("Initialisation de l'orchestrateur")
    print("Initialisation de l'orchestrateur")
    
    orchestrator = Orchestrator(
        output_dir="output/",
        json_agent=json_agent,
        image_sorter=image_sorter,
        image_analyser=image_analyser,
        report_generator=report_generator
    )
    
    # Vérification du ticket spécifique si fourni
    specific_ticket_path = None
    if ticket_id:
        target_ticket = f"ticket_{ticket_id}"
        specific_ticket_path = os.path.join("output", target_ticket)
        
        if not os.path.exists(specific_ticket_path):
            logger.error(f"Le ticket {target_ticket} n'existe pas")
            print(f"ERREUR: Le ticket {target_ticket} n'existe pas")
            return
            
        logger.info(f"Ticket spécifique à traiter: {specific_ticket_path}")
        print(f"Ticket spécifique à traiter: {target_ticket}")
    
    # Exécution de l'orchestrateur
    total_start_time = time.time()
    logger.info("Début de l'exécution de l'orchestrateur")
    print("Début de l'exécution de l'orchestrateur")
    
    orchestrator.executer(specific_ticket_path)
    
    total_time = time.time() - total_start_time
    logger.info(f"Fin de l'exécution de l'orchestrateur (durée: {total_time:.2f} secondes)")
    print(f"Fin de l'exécution de l'orchestrateur (durée: {total_time:.2f} secondes)")

if __name__ == "__main__":
    print("Démarrage du test de l'orchestrateur")
    
    # Vérifier si un ID de ticket est passé en argument
    ticket_id = None
    if len(sys.argv) > 1:
        ticket_id = sys.argv[1]
        print(f"ID de ticket fourni en argument: {ticket_id}")
    
    test_orchestrator(ticket_id)
    print("Test terminé")