from langgraph.graph import StateGraph, MessagesState
from typing import List


class CustomState(MessagesState):
    todo: List[str] # Les tâches en cours, au format JSON

    ragQuery: str # Requête envoyée au RAG, pour le cross-encodeur
    ragDocuments: List[str] # Documents retrouvés par le RAG, pour le cross-encodeur

    lastSummarizedMessage: int # Index du message où l'on s'était arrêté de résumer

    # TODO: Ajouter la source des documents sélectionnés pour la fin du rapport ?


def getState()->StateGraph:
    """
    Retourne un StateGraph prêt à utiliser pour préparer un workflow

    Returns:
        StateGraph: prêt à utiliser
    """
    return StateGraph(CustomState)