Issue with state_modifier in create_react_agent

[luizzappa]

Issue with state_modifier in create_react_agent

I'm facing an issue with the state_modifier in the create_react_agent function, where it doesn't seem to receive the complete state from the graph.

Expected behavior

In the example below, I expected the state to include:

 {
        'messages': [HumanMessage(content='Hi', id='f4b9ebee-5b07-4df9-b408-36900cd5d57b'), 
+       'my_list': [1, 2, 3], 
        'is_last_step': False
}

However, the state that is being passed is incomplete, showing only:

 {
        'messages': [HumanMessage(content='Hi', id='f4b9ebee-5b07-4df9-b408-36900cd5d57b'),
        'is_last_step': False
}

Here’s the basic code illustrating the issue:

class BaseState(TypedDict):
    messages: Annotated[List[BaseMessage], operator.add]
    my_list: List[int]

graph = StateGraph(BaseState)

def _bind_to_system_message(state):
    print(state) # Problem here
    return "system prompt"

agent = create_react_agent(model=model_smart, tools=[], state_modifier=_bind_to_system_message)

graph.add_node('agent', agent)
graph.add_edge(START, 'agent')
graph.add_edge('agent', END)
chain = graph.compile()

for s in chain.invoke({
        "messages": [HumanMessage(content="Hi")],
        "my_list": [1, 2, 3]
    }):
    print(s)
    print('---')

Question

Is this behavior expected for the state_modifier, or is it a bug causing the state to not be passed completely?

[vbarda]

you need to pass state_schema to create_react_agent as well to make sure that agent is aware of the new schema you want to support, just make sure to subclass it from MessagesState. hope this helps!

[luizzappa]

Thank you, it worked!

+ class BaseState(MessagesState):
- class BaseState(TypedDict):
    messages: Annotated[List[BaseMessage], operator.add]
    my_list: List[int]
+   is_last_step: bool

graph = StateGraph(BaseState)

def _bind_to_system_message(state):
    print(state) # Problem here
    return "system prompt"

+ agent = create_react_agent(model=model_smart, tools=[], state_modifier=_bind_to_system_message, state_schema=BaseState)
- agent = create_react_agent(model=model_smart, tools=[], state_modifier=_bind_to_system_message)

graph.add_node('agent', agent)
graph.add_edge(START, 'agent')
graph.add_edge('agent', END)
chain = graph.compile()

for s in chain.invoke({
        "messages": [HumanMessage(content="Hi")],
        "my_list": [1, 2, 3]
    }):
    print(s)
    print('---')