scripts/test_agent.py

#!/usr/bin/env python
"""
Script de test pour l'agent de parsing de leaderboard.
Ce script permet de tester l'agent en standalone en lui donnant une URL.

Usage:
    python test_agent.py <url>
    
Exemple:
    python test_agent.py https://lmarena-ai-chatbot-arena-leaderboard.hf.space
"""
import json
import os
import sys
import argparse
import datetime
from dotenv import load_dotenv

# Assurez-vous que le répertoire parent est dans le chemin d'importation
script_dir = os.path.dirname(os.path.abspath(__file__))
root_dir = os.path.dirname(script_dir)
sys.path.insert(0, root_dir)

from src.agent import get_default_model, process_leaderboard
from src.browser import initialize_driver, close_driver, take_initial_screenshot


def main():
    """
    Fonction principale qui teste l'agent sur une URL spécifiée.
    """
    # Charger les variables d'environnement
    load_dotenv()
    
    # Parse command line arguments
    parser = argparse.ArgumentParser(description="Test de l'agent de parsing de leaderboard")
    parser.add_argument("url", help="URL du leaderboard à parser")
    parser.add_argument("--save", help="Chemin où sauvegarder le résultat JSON (optionnel)")
    parser.add_argument("--uid", help="UID à utiliser pour la capture d'écran (optionnel)")
    parser.add_argument("--wait", type=int, default=10, help="Temps d'attente initial en secondes (défaut: 10)")
    args = parser.parse_args()
    
    # Assurez-vous que nous sommes dans le bon répertoire
    os.chdir(root_dir)
    
    # Vérifier que la clé API est disponible
    if not os.getenv("OPENAI_API_KEY"):
        print("ERREUR: La variable d'environnement OPENAI_API_KEY n'est pas définie.")
        print("Veuillez créer un fichier .env à la racine du projet avec votre clé API.")
        print("Exemple: OPENAI_API_KEY=votre-clé-api")
        sys.exit(1)
    
    # Créer le répertoire d'images si nécessaire
    if args.uid:
        images_dir = os.path.join("data", "images")
        os.makedirs(images_dir, exist_ok=True)
    
    # Obtenir le modèle par défaut
    model = get_default_model()
    
    print(f"Test de l'agent sur l'URL: {args.url}")
    if args.uid:
        print(f"UID utilisé pour la capture d'écran: {args.uid}")
    print(f"Temps d'attente initial: {args.wait} secondes")
    
    # Initialiser le navigateur et prendre une capture d'écran initiale avec le temps d'attente personnalisé
    initialize_driver()
    
    if args.uid:
        # Remplacer la fonction take_initial_screenshot par une version personnalisée avec le temps d'attente spécifié
        import time
        from helium import go_to
        from io import BytesIO
        from PIL import Image
        
        # Naviguer vers l'URL
        go_to(args.url)
        
        # Attendre que la page se charge
        print(f"Attente de {args.wait} secondes pour le chargement complet de la page...")
        time.sleep(args.wait)
        
        # Prendre la capture d'écran
        from src.browser import driver
        png_bytes = driver.get_screenshot_as_png()
        image = Image.open(BytesIO(png_bytes))
        
        # Sauvegarder la capture d'écran
        images_dir = os.path.join("data", "images")
        os.makedirs(images_dir, exist_ok=True)
        screenshot_path = os.path.join(images_dir, f"{args.uid}.png")
        image.save(screenshot_path)
        
        print(f"Capture d'écran initiale sauvegardée dans: {screenshot_path}")
    
    # Fermer le navigateur pour le réinitialiser
    close_driver()
    
    # Traiter le leaderboard
    result = process_leaderboard(args.url, model, 0, args.uid)
    
    # Ajouter des métadonnées
    result["url"] = args.url
    if args.uid:
        result["uid"] = args.uid
        result["screenshot"] = f"images/{args.uid}.png" if os.path.exists(os.path.join("data", "images", f"{args.uid}.png")) else None
    
    # Afficher le résultat en JSON
    json_result = json.dumps(result, indent=2)
    print("\nRésultat JSON:")
    print(json_result)
    
    # Sauvegarder le résultat si demandé
    if args.save:
        with open(args.save, "w") as f:
            f.write(json_result)
        print(f"\nRésultat sauvegardé dans: {args.save}")


if __name__ == "__main__":
    main()