llm_lab/examples/test_mistralapi.py
2025-03-26 15:40:31 +01:00

103 lines
3.4 KiB
Python

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Script d'exemple pour montrer l'utilisation de la classe MistralAPI
"""
import os
import sys
import argparse
from pathlib import Path
# Ajouter le répertoire parent au sys.path pour importer les modules
sys.path.append(str(Path(__file__).parent.parent))
from core.factory import LLMFactory
def afficher_modeles_disponibles():
"""Affiche les modèles disponibles via l'API Mistral"""
# Création d'une instance MistralAPI via la factory
llm = LLMFactory().create("mistralapi")
try:
modeles = llm.obtenir_liste_modeles()
print("\n=== Modèles disponibles sur l'API Mistral ===")
if not modeles:
print("Aucun modèle trouvé ou erreur de récupération.")
return
for i, modele in enumerate(modeles, 1):
print(f"{i}. {modele}")
except Exception as e:
print(f"Erreur lors de la récupération des modèles: {e}")
def generer_reponse(prompt, modele=None, temperature=None, max_tokens=None):
"""Génère une réponse avec MistralAPI"""
# Création d'une instance MistralAPI via la factory
llm = LLMFactory().create("mistralapi")
# Configuration des paramètres si spécifiés
if modele:
llm.modele = modele
if temperature is not None:
llm.parametres_generation["temperature"] = temperature
if max_tokens is not None:
llm.parametres_generation["max_tokens"] = max_tokens
# Affichage des paramètres utilisés
print(f"\n=== Génération avec MistralAPI ===")
print(f"Modèle: {llm.modele}")
print(f"Température: {llm.parametres_generation['temperature']}")
print(f"Max tokens: {llm.parametres_generation.get('max_tokens', 'Non spécifié')}")
print(f"\nPrompt: {prompt}\n")
try:
# Génération de la réponse
reponse = llm.generate(prompt)
# Affichage de la réponse
print("\n=== Réponse ===")
print(reponse)
return reponse
except Exception as e:
print(f"Erreur lors de la génération: {e}")
return None
if __name__ == "__main__":
parser = argparse.ArgumentParser(description="Test de l'API Mistral")
parser.add_argument("--list-models", action="store_true", help="Afficher les modèles disponibles")
parser.add_argument("--prompt", help="Prompt pour générer une réponse")
parser.add_argument("--model", help="Nom du modèle à utiliser (ex: mistral-large-latest)")
parser.add_argument("--temperature", type=float, help="Température pour la génération (0.0-1.0)")
parser.add_argument("--max-tokens", type=int, help="Nombre maximum de tokens à générer")
args = parser.parse_args()
# Vérification de la clé API dans l'environnement
if not os.environ.get("MISTRAL_API_KEY"):
print("Erreur: Variable d'environnement MISTRAL_API_KEY non définie.")
print("Veuillez définir cette variable avec votre clé API Mistral.")
print("Exemple: export MISTRAL_API_KEY=votre_clé_api")
sys.exit(1)
# Exécution en fonction des arguments
if args.list_models:
afficher_modeles_disponibles()
elif args.prompt:
generer_reponse(
args.prompt,
modele=args.model,
temperature=args.temperature,
max_tokens=args.max_tokens
)
else:
parser.print_help()