From 97a2ce9d1a32190db051a905ccaa873db081add7 Mon Sep 17 00:00:00 2001 From: vbarda Date: Thu, 23 Jan 2025 15:43:42 -0500 Subject: [PATCH] another test --- libs/langgraph/tests/test_prebuilt.py | 17 +++++++++++++++++ 1 file changed, 17 insertions(+) diff --git a/libs/langgraph/tests/test_prebuilt.py b/libs/langgraph/tests/test_prebuilt.py index 2675ba664b..9db0348e66 100644 --- a/libs/langgraph/tests/test_prebuilt.py +++ b/libs/langgraph/tests/test_prebuilt.py @@ -275,6 +275,23 @@ def prompt(state): assert response == expected_response +async def test_callable_prompt_async(): + async def prompt(state): + modified_message = f"Bar {state['messages'][-1].content}" + return [HumanMessage(content=modified_message)] + + for agent in ( + create_react_agent(FakeToolCallingModel(), [], prompt=prompt), + create_react_agent(FakeToolCallingModel(), [], state_modifier=prompt), + ): + inputs = [HumanMessage("hi?")] + response = await agent.ainvoke({"messages": inputs}) + expected_response = { + "messages": inputs + [AIMessage(content="Bar hi?", id="0")] + } + assert response == expected_response + + def test_runnable_prompt(): messages_modifier = RunnableLambda( lambda messages: [HumanMessage(content=f"Baz {messages[-1].content}")]