core/state.py

from langgraph.graph import MessagesState
import operator
from typing_extensions import TypedDict, Annotated, List, Sequence
from langchain_core.messages import BaseMessage
from langgraph.graph.message import add_messages


class State(MessagesState):
    summary: str
    question: str
    chunked_last_tool_call: bool
    file_reference: str  # Attachment file reference: a path, URL, or unique ID


class MathAgentState(MessagesState):
    """
    State for the math agent containing message history and research metadata.

    This state tracks the agent's conversation, iteration count for limiting
    tool calls, the research topic being investigated, compressed findings,
    and raw research notes for detailed analysis.
    """
    tool_call_iterations: int
    question: str
    compressed_research: str
    raw_notes: Annotated[List[str], operator.add]


class MathAgentOutputState(MessagesState):
    """
    Output state for the math agent containing final results.

    This represents the final output of the solving process with steps
    and raw notes from the solving process.
    """
    compressed_research: str
    raw_notes: Annotated[List[str], operator.add]