llm_ticket3/test_image_sorter.py

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

import os
import sys
import argparse
import json
from llm_classes.llama_vision import LlamaVision
from agents.llama_vision.agent_image_sorter import AgentImageSorter

def parse_args():
    parser = argparse.ArgumentParser(description="Test de l'agent de tri d'images")
    parser.add_argument("image_path", help="Chemin vers l'image à analyser")
    parser.add_argument("--debug", "-d", action="store_true", help="Mode debug")
    parser.add_argument("--ticket_id", "-t", default="T9999", help="ID du ticket pour les tests")
    return parser.parse_args()

def main():
    args = parse_args()
    image_path = args.image_path
    
    if not os.path.exists(image_path):
        print(f"Erreur: L'image {image_path} n'existe pas")
        return 1
        
    print(f"=== TEST DE L'AGENT DE TRI D'IMAGES ===")
    print(f"Image: {image_path}")
    
    # Initialiser le modèle LlamaVision
    try:
        print("Initialisation du modèle LlamaVision...")
        llm = LlamaVision()  # modèle par défaut
        print(f"Modèle initialisé: {llm.modele}")
    except Exception as e:
        print(f"Erreur lors de l'initialisation du modèle: {e}")
        return 1
        
    # Initialiser l'agent de tri
    try:
        print("Initialisation de l'agent de tri d'images...")
        agent = AgentImageSorter(llm)
        print("Agent initialisé")
    except Exception as e:
        print(f"Erreur lors de l'initialisation de l'agent: {e}")
        return 1
        
    # Exécuter l'agent sur l'image
    try:
        print("\nExécution de l'agent sur l'image...")
        result = agent.executer(image_path)
        
        # Afficher le résultat
        print("\n=== RÉSULTAT ===")
        print(f"Image pertinente: {result['is_relevant']}")
        print(f"Raison: {result['reason']}")
        
        if args.debug:
            print("\n=== RÉPONSE BRUTE ===")
            print(result['raw_response'])
            print("\n=== OCR ===")
            print(f"OCR FR: {result['ocr_fr'] or 'Aucun texte détecté'}")
            print(f"OCR EN: {result['ocr_en'] or 'Aucune traduction'}")
        
        # Sauvegarder le résultat dans un fichier JSON
        output_file = f"test_tri_{os.path.basename(image_path)}.json"
        with open(output_file, "w", encoding="utf-8") as f:
            json.dump(result, f, ensure_ascii=False, indent=2)
        print(f"\nRésultat enregistré dans {output_file}")
        
        return 0
    except Exception as e:
        print(f"Erreur lors de l'exécution de l'agent: {e}")
        import traceback
        traceback.print_exc()
        return 1

if __name__ == "__main__":
    sys.exit(main())