Spaces:
Build error
Build error
File size: 4,273 Bytes
0821095 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 |
#!/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() |