diff --git a/libs/langgraph/tests/test_pregel_async.py b/libs/langgraph/tests/test_pregel_async.py index 4995addb6..f59433a07 100644 --- a/libs/langgraph/tests/test_pregel_async.py +++ b/libs/langgraph/tests/test_pregel_async.py @@ -21,6 +21,9 @@ from uuid import UUID import httpx import pytest +from langchain_core.messages import ( + ToolCall, +) from langchain_core.runnables import ( RunnableConfig, RunnableLambda, @@ -3877,6 +3880,12 @@ async def test_prebuilt_tool_chat() -> None: ] +# defined outside to allow deserializer to see it +class ToolInput(BaseModel, arbitrary_types_allowed=True): + call: ToolCall + my_session: httpx.AsyncClient + + async def test_state_graph_packets() -> None: from langchain_core.language_models.fake_chat_models import ( FakeMessagesListChatModel, @@ -3885,7 +3894,6 @@ async def test_state_graph_packets() -> None: AIMessage, BaseMessage, HumanMessage, - ToolCall, ToolMessage, ) from langchain_core.tools import tool @@ -3941,19 +3949,15 @@ async def test_state_graph_packets() -> None: # Logic to decide whether to continue in the loop or exit if tool_calls := data["messages"][-1].tool_calls: return [ - Send("tools", {"call": tool_call, "my_session": data["session"]}) + Send("tools", ToolInput(call=tool_call, my_session=data["session"])) for tool_call in tool_calls ] else: return END - class ToolInput(TypedDict): - call: ToolCall - my_session: httpx.Client - async def tools_node(input: ToolInput, config: RunnableConfig) -> AgentState: - assert isinstance(input["my_session"], httpx.AsyncClient) - tool_call = input["call"] + assert isinstance(input.my_session, httpx.AsyncClient) + tool_call = input.call await asyncio.sleep(tool_call["args"].get("idx", 0) / 10) output = await tools_by_name[tool_call["name"]].ainvoke( tool_call["args"], config