from typing import List, Dict, TypedDict, Sequence, Union, Annotated

from langchain_core.messages import BaseMessage
from langgraph.graph.message import add_messages


class Context(TypedDict):
    """PCS + geonames context payload for common tasks like recommendations.
    """
    subject: List[str]
    population: List[str]
    geography: List[Union[str, int]]


class AgentState(TypedDict):
    """State of the chat agent for the execution graph(s).
    """
    # The add_messages function defines how an update should be processed
    # Default is to replace. add_messages says "append"
    messages: Annotated[Sequence[BaseMessage], add_messages]
    user_input: str
    org_dict: Dict

    # Recommendation-specific fields
    intent: str
    context: Context
    recommendation: str