import time
import asyncio
from langchain_core.messages import AIMessage, HumanMessage
from src.agents.main_agent import multi_agent_graph

async def process_query(user_input):
    messages = [HumanMessage(content=user_input)]
    config = {"configurable": {"thread_id": "1"}}

    # Mesurer le temps d'exécution
    start_time = time.time()
    
    result = await multi_agent_graph.ainvoke({"messages": messages}, config=config)
    
    end_time = time.time()
    execution_time = end_time - start_time  # Temps écoulé

    # Afficher les messages de réponse
    #for message in result["messages"]:
    #    message.pretty_print()
    
    print(result["messages"][-1].content)

    print(f"Temps d'exécution: {execution_time:.4f} secondes")

async def main():
    print("Bienvenue dans le système de requêtes. Tapez 'exit' pour quitter.")

    while True:
        # Demander à l'utilisateur de saisir une requête
        user_input = input("Veuillez saisir votre requête : ")
        
        if user_input.lower() == 'exit':
            print("Au revoir!")
            break
        
        await process_query(user_input)

# Exécuter la boucle d'événements
if __name__ == "__main__":
    asyncio.run(main())