Proposal: optimize syntax expression of AgentState of LangGraph

[Undertone0809]

Checked

• I searched existing ideas and did not find a similar one
• I added a very descriptive title
• I've clearly described the feature request and motivation for it

Feature request

Currently, AgentState is just a normal dict, which means that there are no attribute restrictions, and if the system is complex, each node returns a different field, which is difficult to maintain.

An AgentState like this:

class AgentState(TypedDict):
    # The input string
    input: str
    # The list of previous messages in the conversation
    chat_history: list[BaseMessage]
    # The outcome of a given call to the agent
    # Needs `None` as a valid type, since this is what this will start as
    agent_outcome: Union[AgentAction, AgentFinish, None]
    # List of actions and corresponding observations
    # Here we annotate this with `operator.add` to indicate that operations to
    # this state should be ADDED to the existing values (not overwrite it)
    intermediate_steps: Annotated[list[tuple[AgentAction, str]], operator.add]

Refer to the implementation of AgentProtocol. In AgentProtocol, each step is processed by passing a step into the function. After changing the step, a step is returned.

from agent_protocol import Agent, Step, Task


async def plan(step: Step) -> Step:
    task = await Agent.db.get_task(step.task_id)
    steps = generate_steps(task.input)

    last_step = steps[-1]
    for step in steps[:-1]:
       await Agent.db.create_step(task_id=task.task_id, name=step, ...)

    await Agent.db.create_step(task_id=task.task_id, name=last_step, is_last=True)
    step.output = steps
    return step


async def execute(step: Step) -> Step:
    # Use tools, websearch, etc.
    ...

async def task_handler(task: Task) -> None:
    await Agent.db.create_step(task_id=task.task_id, name="plan", ...)


async def step_handler(step: Step) -> Step:
    if step.name == "plan":
        await plan(step)
    else:
        await execute(step)

   return step


Agent.setup_agent(task_handler, step_handler).start()

In the Step, the AgentProtocol can store some intermediate data of the step, which can be equivalent to the AgentState of each node in the LangGraph. According to this notation, we can define a node function in the following way:

def execute_tools(state: AgentState) -> AgentState:
    ...
    return state

Motivation

Currently, AgentState is just a normal dict, which means that there are no attribute restrictions, and if the system is complex, each node returns a different field, which is difficult to maintain.

Proposal (If applicable)

The proposal is to create a class that inherits from dict, and that has a fixed set of attributes, and that the system can only return these attributes.

Original Usage:

from langchain_core.agents import AgentFinish
from langgraph.prebuilt.tool_executor import ToolExecutor

# This a helper class we have that is useful for running tools
# It takes in an agent action and calls that tool and returns the result
tool_executor = ToolExecutor(tools)


# Define the agent
def run_agent(data):
    agent_outcome = agent_runnable.invoke(data)
    return {"agent_outcome": agent_outcome}


# Define the function to execute tools
def execute_tools(data):
    # Get the most recent agent_outcome - this is the key added in the `agent` above
    agent_action = data["agent_outcome"]
    output = tool_executor.invoke(agent_action)
    return {"intermediate_steps": [(agent_action, str(output))]}


# Define logic that will be used to determine which conditional edge to go down
def should_continue(data):
    # If the agent outcome is an AgentFinish, then we return `exit` string
    # This will be used when setting up the graph to define the flow
    if isinstance(data["agent_outcome"], AgentFinish):
        return "end"
    # Otherwise, an AgentAction is returned
    # Here we return `continue` string
    # This will be used when setting up the graph to define the flow
    else:
        return "continue"

Optimized Usage:

1. Use AgentState class to define data, inherit from pydantic's BaseModel, which allows developers to define the description of data using Field when defining AgentState.
2. 2. For general nodes, return an AgentState class instead of a dict, which allows developers to better maintain data.

from typing import TypedDict, Annotated, List, Union
from langchain_core.agents import AgentAction, AgentFinish
from langchain_core.messages import BaseMessage
from enum import Enum
from pydantic import BaseModel, Field
from langchain import hub
from langchain.agents import create_openai_functions_agent
from langchain_openai.chat_models import ChatOpenAI
from langchain_community.tools.tavily_search import TavilySearchResults
from enum import Enum
from langchain_core.agents import AgentFinish
from langgraph.prebuilt.tool_executor import ToolExecutor
from langgraph.graph import END, StateGraph


class BaseAgentState(BaseModel):
    ...


class AgentState(BaseAgentState):
    input: str = Field(..., description="The input string")
    chat_history: List[BaseMessage] = Field(
        [], description="The list of previous messages in the conversation"
    )
    agent_outcome: Union[AgentAction, AgentFinish, None] = Field(
        None, description="The outcome of a given call to the agent"
    )
    intermediate_steps: List[tuple[AgentAction, str]] = Field(
        [], description="List of actions and corresponding observations"
    )


class Condition(str, Enum):
    CONTINUE = "continue"
    END = "end"


tools = [TavilySearchResults(max_results=1)]

# Get the prompt to use - you can modify this!
prompt = hub.pull("hwchase17/openai-functions-agent")

# Choose the LLM that will drive the agent
llm = ChatOpenAI(model="gpt-3.5-turbo-1106", streaming=True)

# Construct the OpenAI Functions agent
agent_runnable = create_openai_functions_agent(llm, tools, prompt)

# This a helper class we have that is useful for running tools
# It takes in an agent action and calls that tool and returns the result
tool_executor = ToolExecutor(tools)


# Define the agent
def run_agent(state: AgentState) -> AgentState:
    agent_outcome = agent_runnable.invoke(state)
    state.agent_outcome = agent_outcome
    return state


# Define the function to execute tools
def execute_tools(state: AgentState) -> AgentState:
    # Get the most recent agent_outcome - this is the key added in the `agent` above
    agent_action = state.agent_outcome
    output = tool_executor.invoke(agent_action)
    state.intermediate_steps = [(agent_action, str(output))]
    return state


# Define logic that will be used to determine which conditional edge to go down
def should_continue(state: AgentState) -> Condition:
    # If the agent outcome is an AgentFinish, then we return `exit` string
    # This will be used when setting up the graph to define the flow
    if isinstance(state.agent_outcome, AgentFinish):
        return Condition.CONTINUE
    # Otherwise, an AgentAction is returned
    # Here we return `continue` string
    # This will be used when setting up the graph to define the flow
    else:
        return Condition.END

class NodeType(str, Enum):
    AGENT = "agent"
    ACTION = "action"


# Define a new graph
workflow = StateGraph(AgentState)

# Define the two nodes we will cycle between
workflow.add_node(NodeType.AGENT, run_agent)
workflow.add_node(NodeType.ACTION, execute_tools)

# Set the entrypoint as `agent`
# This means that this node is the first one called
workflow.set_entry_point(NodeType.AGENT)

# We now add a conditional edge
workflow.add_conditional_edges(
    # First, we define the start node. We use `agent`.
    # This means these are the edges taken after the `agent` node is called.
    NodeType.AGENT,
    # Next, we pass in the function that will determine which node is called next.
    should_continue,
    # Finally we pass in a mapping.
    # The keys are strings, and the values are other nodes.
    # END is a special node marking that the graph should finish.
    # What will happen is we will call `should_continue`, and then the output of that
    # will be matched against the keys in this mapping.
    # Based on which one it matches, that node will then be called.
    {
        # If `tools`, then we call the tool node.
        Condition.CONTINUE: NodeType.ACTION,
        # Otherwise we finish.
        Condition.END: END,
    },
)

# We now add a normal edge from `tools` to `agent`.
# This means that after `tools` is called, `agent` node is called next.
workflow.add_edge(NodeType.ACTION, NodeType.AGENT)

# Finally, we compile it!
# This compiles it into a LangChain Runnable,
# meaning you can use it as you would any other runnable
app = workflow.compile()

Other

1. Encourage the use of enumeration to describe node names, state names, to avoid the misspelling of words caused by low-level errors, in the development of complex systems, data declaration is particularly important, if you agree with this development paradigm, I can help modify the examples of these declarations to enum type.
2. This is only a draft, if you have any suggestions, please let me know. If you need my help, I will be happy to develop and follow up.

[Undertone0809]

@nfcampos @hinthornw @hwaking @rlancemartin @baskaryan Hey, how about this suggestion, looking forward to your reply~

[nfcampos]

hi, thanks for this, this is actually already possible today, you can already use a pydantic base model as your state object, see example here https://github.com/langchain-ai/langgraph/blob/main/examples/state-model.ipynb