File size: 2,283 Bytes
c2df2b9 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 | """Define the state structures for the agent."""
from __future__ import annotations
from dataclasses import dataclass, field
from typing import Sequence
from langchain_core.messages import AnyMessage
from langgraph.graph import add_messages
from langgraph.managed import IsLastStep
from typing_extensions import Annotated
@dataclass
class InputState:
"""Defines the input state for the agent, representing a narrower interface to the outside world.
This class is used to define the initial state and structure of incoming data.
"""
messages: Annotated[Sequence[AnyMessage], add_messages] = field(
default_factory=list
)
"""
Messages tracking the primary execution state of the agent.
Typically accumulates a pattern of:
1. HumanMessage - user input
2. AIMessage with .tool_calls - agent picking tool(s) to use to collect information
3. ToolMessage(s) - the responses (or errors) from the executed tools
4. AIMessage without .tool_calls - agent responding in unstructured format to the user
5. HumanMessage - user responds with the next conversational turn
Steps 2-5 may repeat as needed.
The `add_messages` annotation ensures that new messages are merged with existing ones,
updating by ID to maintain an "append-only" state unless a message with the same ID is provided.
"""
@dataclass
class State(InputState):
"""Represents the complete state of the agent, extending InputState with additional attributes.
This class can be used to store any information needed throughout the agent's lifecycle.
"""
is_last_step: IsLastStep = field(default=False)
"""
Indicates whether the current step is the last one before the graph raises an error.
This is a 'managed' variable, controlled by the state machine rather than user code.
It is set to 'True' when the step count reaches recursion_limit - 1.
"""
# Additional attributes can be added here as needed.
# Common examples include:
# retrieved_documents: List[Document] = field(default_factory=list)
# extracted_entities: Dict[str, Any] = field(default_factory=dict)
# api_connections: Dict[str, Any] = field(default_factory=dict)
|