from typing import List, Optional, TypedDict

from langchain_core.messages import AIMessage, HumanMessage


class AgentState(TypedDict):
    """
    Represents the state of an agent during a multi-step query process.
    Attributes:
        original_query (str): The initial query provided by the user.
        subqueries (List[str]): A list of subqueries derived from the
                                original query.
        current_step (int): The index of the current step or subquery
                            being processed.
        results (List[str]): The results obtained for each subquery.
        final_answer (Optional[str]): The final answer generated by the agent,
                                      if available.
        messages (List[HumanMessage | AIMessage]): The conversation history,
                                                   including both human and AI
                                                   messages.
    """
    original_query: str
    subqueries: List[str]
    current_step: int
    results: List[str]
    final_answer: Optional[str]
    messages: List[HumanMessage | AIMessage]