From 9e066554ba7d2eb54e3f2695e955e1488dda2cf9 Mon Sep 17 00:00:00 2001
From: Vadym Barda <vadym@langchain.dev>
Date: Thu, 23 Jan 2025 09:29:25 -0500
Subject: [PATCH] langgraph: allow async state modifier in create_react_agent
 (#3161)

Fixes https://github.com/langchain-ai/langgraph/issues/2875
---
 .../langgraph/prebuilt/chat_agent_executor.py |  7 +++
 libs/langgraph/tests/test_prebuilt.py         | 44 +++++++++++++++++++
 2 files changed, 51 insertions(+)

diff --git a/libs/langgraph/langgraph/prebuilt/chat_agent_executor.py b/libs/langgraph/langgraph/prebuilt/chat_agent_executor.py
index ceeba34b3..005ae59d3 100644
--- a/libs/langgraph/langgraph/prebuilt/chat_agent_executor.py
+++ b/libs/langgraph/langgraph/prebuilt/chat_agent_executor.py
@@ -1,3 +1,4 @@
+import inspect
 from typing import (
     Callable,
     Literal,
@@ -91,6 +92,12 @@ def _get_state_modifier_runnable(
             lambda state: [state_modifier] + state["messages"],
             name=STATE_MODIFIER_RUNNABLE_NAME,
         )
+    elif inspect.iscoroutinefunction(state_modifier):
+        state_modifier_runnable = RunnableCallable(
+            None,
+            state_modifier,
+            name=STATE_MODIFIER_RUNNABLE_NAME,
+        )
     elif callable(state_modifier):
         state_modifier_runnable = RunnableCallable(
             state_modifier,
diff --git a/libs/langgraph/tests/test_prebuilt.py b/libs/langgraph/tests/test_prebuilt.py
index 99263062a..04ce3ee69 100644
--- a/libs/langgraph/tests/test_prebuilt.py
+++ b/libs/langgraph/tests/test_prebuilt.py
@@ -346,6 +346,50 @@ def test_state_modifier_with_store():
     assert response["messages"][-1].content == "foo-hi"
 
 
+async def test_state_modifier_with_store_async():
+    async def add(a: int, b: int):
+        """Adds a and b"""
+        return a + b
+
+    in_memory_store = InMemoryStore()
+    await in_memory_store.aput(
+        ("memories", "1"), "user_name", {"data": "User name is Alice"}
+    )
+    await in_memory_store.aput(
+        ("memories", "2"), "user_name", {"data": "User name is Bob"}
+    )
+
+    async def modify(state, config, *, store):
+        user_id = config["configurable"]["user_id"]
+        system_str = (await store.aget(("memories", user_id), "user_name")).value[
+            "data"
+        ]
+        return [SystemMessage(system_str)] + state["messages"]
+
+    async def modify_no_store(state, config):
+        return SystemMessage("foo") + state["messages"]
+
+    model = FakeToolCallingModel()
+
+    # test state modifier that uses store works
+    agent = create_react_agent(
+        model, [add], state_modifier=modify, store=in_memory_store
+    )
+    response = await agent.ainvoke(
+        {"messages": [("user", "hi")]}, {"configurable": {"user_id": "1"}}
+    )
+    assert response["messages"][-1].content == "User name is Alice-hi"
+
+    # test state modifier that doesn't use store works
+    agent = create_react_agent(
+        model, [add], state_modifier=modify_no_store, store=in_memory_store
+    )
+    response = await agent.ainvoke(
+        {"messages": [("user", "hi")]}, {"configurable": {"user_id": "2"}}
+    )
+    assert response["messages"][-1].content == "foo-hi"
+
+
 @pytest.mark.parametrize("tool_style", ["openai", "anthropic"])
 def test_model_with_tools(tool_style: str):
     model = FakeToolCallingModel(tool_style=tool_style)