from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
from langchain_core.messages import HumanMessage
from src.agents.main_agent import multi_agent_graph
from src.agents.rag_agent import invocation_state

import time

invocation_state.invocation_type = 'api'

app = FastAPI()

# Modèle Pydantic pour valider l'entrée de l'utilisateur
class UserInput(BaseModel):
    question: str

@app.post("/ask")
async def ask_question(user_input: UserInput):
    messages = [HumanMessage(content=user_input.question)]
    config = {"configurable": {"thread_id": "1"}}

    # Mesurer le temps d'exécution
    start_time = time.time()
    
    try:
        result = await multi_agent_graph.ainvoke({"messages": messages}, config=config)
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))
    
    end_time = time.time()
    execution_time = end_time - start_time  # Temps écoulé

    # Retourner la réponse
    response_message = result["messages"][-1].content
    print(f"Temps d'exécution: {execution_time:.4f} secondes")
    return {"response": response_message, "execution_time": execution_time}