from typing import Annotated, NotRequired, TypedDict
from langgraph.graph import add_messages
from langgraph.managed.is_last_step import RemainingSteps


class State(TypedDict):
    messages: Annotated[list, add_messages]
    remaining_steps: RemainingSteps
    run_context: NotRequired[dict]
    intent: NotRequired[dict]
    task_plan: NotRequired[dict]
    action_plan: NotRequired[dict]
    action_states: NotRequired[list[dict]]
    identity_resolution: NotRequired[dict]
    tool_trace: NotRequired[list[dict]]
    observations: NotRequired[dict]
    artifacts: NotRequired[dict]
    scientific_ledger: NotRequired[dict]
    validation: NotRequired[dict]
    tool_guard: NotRequired[dict]
    clarification: NotRequired[dict]
    memory_refs: NotRequired[dict]
    final_output: NotRequired[dict]
    completion_validation: NotRequired[dict]
    validation_attempts: NotRequired[int]


class MultiAgentState(TypedDict):
    question: str
    first_router_response: Annotated[list, add_messages]
    regular_response: Annotated[list, add_messages]
    feedback_response: Annotated[list, add_messages]
    geometry_response: Annotated[list, add_messages]
    parameter_response: Annotated[list, add_messages]
    opt_response: Annotated[list, add_messages]
    end_response: Annotated[list, add_messages]