From 834fa8932fd8be03245be6cb178372870dae744b Mon Sep 17 00:00:00 2001 From: Sydney Runkle Date: Wed, 21 Jan 2026 12:43:32 -0500 Subject: [PATCH] async --- .../langgraph/prebuilt/chat_agent_executor.py | 26 +++- libs/prebuilt/langgraph/prebuilt/tool_node.py | 46 +++++- libs/prebuilt/tests/test_tool_node.py | 138 ++++++++++++++++++ 3 files changed, 201 insertions(+), 9 deletions(-) diff --git a/libs/prebuilt/langgraph/prebuilt/chat_agent_executor.py b/libs/prebuilt/langgraph/prebuilt/chat_agent_executor.py index 4076cfa4b..1305f3098 100644 --- a/libs/prebuilt/langgraph/prebuilt/chat_agent_executor.py +++ b/libs/prebuilt/langgraph/prebuilt/chat_agent_executor.py @@ -287,7 +287,10 @@ def create_react_agent( [StateSchema, Runtime[ContextT]], Awaitable[Runnable[LanguageModelInput, BaseMessage]], ], - tools: Sequence[BaseTool | Callable | dict[str, Any]] | ToolNode, + tools: Sequence[BaseTool | Callable | dict[str, Any]] + | Callable[[], Sequence[BaseTool]] + | Callable[[], Awaitable[Sequence[BaseTool]]] + | ToolNode, *, prompt: Prompt | None = None, response_format: StructuredResponseSchema @@ -355,8 +358,9 @@ def create_react_agent( `.bind_tools()` and support required functionality. Bound tools must be a subset of those specified in the `tools` parameter. - tools: A list of tools or a `ToolNode` instance. + tools: A list of tools, a `ToolNode` instance, or a callable that returns tools. If an empty list is provided, the agent will consist of a single LLM node without tool calling. + Callable tools providers (both sync and async) enable dynamic tool selection at runtime. prompt: An optional prompt for the LLM. Can take a few different forms: - `str`: This is converted to a `SystemMessage` and added to the beginning of the list of messages in `state["messages"]`. @@ -546,18 +550,30 @@ def create_react_agent( ) llm_builtin_tools: list[dict] = [] + is_dynamic_tools = callable(tools) and not isinstance(tools, ToolNode) if isinstance(tools, ToolNode): tool_classes = list(tools.tools_by_name.values()) tool_node = tools + elif is_dynamic_tools: + # Dynamic tools provider - pass directly to ToolNode + tool_node = ToolNode( + cast( + "Callable[[], Sequence[BaseTool]] | Callable[[], Awaitable[Sequence[BaseTool]]]", + tools, + ) + ) + # For dynamic tools, we can't know the tools at compile time + tool_classes = [] else: - llm_builtin_tools = [t for t in tools if isinstance(t, dict)] - tool_node = ToolNode([t for t in tools if not isinstance(t, dict)]) + tools_seq = cast("Sequence[BaseTool | Callable | dict[str, Any]]", tools) + llm_builtin_tools = [t for t in tools_seq if isinstance(t, dict)] + tool_node = ToolNode([t for t in tools_seq if not isinstance(t, dict)]) tool_classes = list(tool_node.tools_by_name.values()) is_dynamic_model = not isinstance(model, (str, Runnable)) and callable(model) is_async_dynamic_model = is_dynamic_model and inspect.iscoroutinefunction(model) - tool_calling_enabled = len(tool_classes) > 0 + tool_calling_enabled = len(tool_classes) > 0 or is_dynamic_tools if not is_dynamic_model: if isinstance(model, str): diff --git a/libs/prebuilt/langgraph/prebuilt/tool_node.py b/libs/prebuilt/langgraph/prebuilt/tool_node.py index e7970081b..cce131084 100644 --- a/libs/prebuilt/langgraph/prebuilt/tool_node.py +++ b/libs/prebuilt/langgraph/prebuilt/tool_node.py @@ -738,7 +738,9 @@ class ToolNode(RunnableCallable): def __init__( self, - tools: Sequence[BaseTool | Callable] | Callable[[], Sequence[BaseTool]], + tools: Sequence[BaseTool | Callable] + | Callable[[], Sequence[BaseTool]] + | Callable[[], Awaitable[Sequence[BaseTool]]], *, name: str = "tools", tags: list[str] | None = None, @@ -773,11 +775,19 @@ class ToolNode(RunnableCallable): self._awrap_tool_call = awrap_tool_call self._tools_provider: Callable[[], Sequence[BaseTool]] | None = None + self._async_tools_provider: ( + Callable[[], Awaitable[Sequence[BaseTool]]] | None + ) = None self._tools_by_name: dict[str, BaseTool] = {} if callable(tools) and not isinstance(tools, (list, tuple)): # It's a dynamic tools provider - self._tools_provider = tools + if inspect.iscoroutinefunction(tools): + self._async_tools_provider = cast( + "Callable[[], Awaitable[Sequence[BaseTool]]]", tools + ) + else: + self._tools_provider = cast("Callable[[], Sequence[BaseTool]]", tools) else: # It's a sequence of tools - process them statically self._tools_by_name = self._build_tools_mapping(tools) @@ -823,7 +833,33 @@ class ToolNode(RunnableCallable): Returns: Dictionary mapping tool names to BaseTool instances. + + Raises: + TypeError: If an async tools provider is used in synchronous context. """ + if self._async_tools_provider is not None: + msg = ( + "Cannot use async tools provider in synchronous context. " + "Use ainvoke() instead of invoke()." + ) + raise TypeError(msg) + if self._tools_provider is not None: + tools = self._tools_provider() + return self._build_tools_mapping(tools, convert_callables=False) + return self._tools_by_name + + async def _aget_tools(self) -> dict[str, BaseTool]: + """Get the current tools mapping asynchronously. + + If an async or sync tools provider was configured, calls it to get + the current tools. Otherwise, returns the statically configured tools. + + Returns: + Dictionary mapping tool names to BaseTool instances. + """ + if self._async_tools_provider is not None: + tools = await self._async_tools_provider() + return self._build_tools_mapping(tools, convert_callables=False) if self._tools_provider is not None: tools = self._tools_provider() return self._build_tools_mapping(tools, convert_callables=False) @@ -833,8 +869,10 @@ class ToolNode(RunnableCallable): def tools_by_name(self) -> dict[str, BaseTool]: """Mapping from tool name to BaseTool instance. - Note: If a dynamic tools provider was configured, this property + Note: If a sync dynamic tools provider was configured, this property calls the provider to get the current tools on each access. + If an async tools provider was configured, this property will raise + a TypeError - use ainvoke() instead. """ return self._get_tools() @@ -889,7 +927,7 @@ class ToolNode(RunnableCallable): config_list = get_config_list(config, len(tool_calls)) # Get tools once at the start of invocation (supports dynamic tools) - tools_by_name = self._get_tools() + tools_by_name = await self._aget_tools() # Construct ToolRuntime instances at the top level for each tool call tool_runtimes = [] diff --git a/libs/prebuilt/tests/test_tool_node.py b/libs/prebuilt/tests/test_tool_node.py index 6cdc23466..6ba0df206 100644 --- a/libs/prebuilt/tests/test_tool_node.py +++ b/libs/prebuilt/tests/test_tool_node.py @@ -2083,3 +2083,141 @@ def test_tool_node_dynamic_tools_type_error() -> None: }, config=_create_config_with_runtime(), ) + + +async def test_tool_node_async_tools_provider() -> None: + """Test ToolNode with an async tools provider callable.""" + + @dec_tool + def add(a: int, b: int) -> int: + """Add two numbers.""" + return a + b + + @dec_tool + def multiply(a: int, b: int) -> int: + """Multiply two numbers.""" + return a * b + + @dec_tool + def subtract(a: int, b: int) -> int: + """Subtract two numbers.""" + return a - b + + # Track which tools are available + available_tools: list[BaseTool] = [add, multiply] + + async def get_tools_async() -> list[BaseTool]: + # Simulate async operation (e.g., fetching tools from a database) + return available_tools + + # Create ToolNode with async dynamic tools provider + tool_node = ToolNode(get_tools_async) + + # Test invoking a tool asynchronously + result = await tool_node.ainvoke( + { + "messages": [ + AIMessage( + "test", + tool_calls=[ + {"name": "add", "args": {"a": 2, "b": 3}, "id": "call_1"} + ], + ) + ] + }, + config=_create_config_with_runtime(), + ) + tool_message = result["messages"][-1] + assert tool_message.content == "5" + + # Test invoking another tool + result = await tool_node.ainvoke( + { + "messages": [ + AIMessage( + "test", + tool_calls=[ + {"name": "multiply", "args": {"a": 4, "b": 5}, "id": "call_2"} + ], + ) + ] + }, + config=_create_config_with_runtime(), + ) + tool_message = result["messages"][-1] + assert tool_message.content == "20" + + # Change the available tools dynamically + available_tools.clear() + available_tools.extend([subtract]) + + # Test that the new tool works + result = await tool_node.ainvoke( + { + "messages": [ + AIMessage( + "test", + tool_calls=[ + {"name": "subtract", "args": {"a": 10, "b": 3}, "id": "call_4"} + ], + ) + ] + }, + config=_create_config_with_runtime(), + ) + tool_message = result["messages"][-1] + assert tool_message.content == "7" + + +def test_tool_node_async_tools_provider_sync_context_error() -> None: + """Test that async tools provider raises TypeError in sync context.""" + + @dec_tool + def add(a: int, b: int) -> int: + """Add two numbers.""" + return a + b + + async def get_tools_async() -> list[BaseTool]: + return [add] + + tool_node = ToolNode(get_tools_async) + + # Should raise TypeError when trying to invoke synchronously + with pytest.raises( + TypeError, + match="Cannot use async tools provider in synchronous context", + ): + tool_node.invoke( + { + "messages": [ + AIMessage( + "test", + tool_calls=[ + {"name": "add", "args": {"a": 2, "b": 3}, "id": "call_1"} + ], + ) + ] + }, + config=_create_config_with_runtime(), + ) + + +def test_tool_node_async_tools_provider_tools_by_name_error() -> None: + """Test that tools_by_name raises TypeError with async tools provider.""" + + @dec_tool + def add(a: int, b: int) -> int: + """Add two numbers.""" + return a + b + + async def get_tools_async() -> list[BaseTool]: + return [add] + + tool_node = ToolNode(get_tools_async) + + # Should raise TypeError when accessing tools_by_name + with pytest.raises( + TypeError, + match="Cannot use async tools provider in synchronous context", + ): + _ = tool_node.tools_by_name