From 2f8ae4fdb4842576647e773890fc4b225419fc46 Mon Sep 17 00:00:00 2001
From: Nuno Campos <nuno@langchain.dev>
Date: Tue, 2 Apr 2024 14:46:14 -0700
Subject: [PATCH] Simplify tool executor

---
 langgraph/prebuilt/tool_executor.py | 31 +++++++++++------------------
 tests/test_pregel.py                |  8 ++++----
 2 files changed, 16 insertions(+), 23 deletions(-)

diff --git a/langgraph/prebuilt/tool_executor.py b/langgraph/prebuilt/tool_executor.py
index b4d86240e..d838579b2 100644
--- a/langgraph/prebuilt/tool_executor.py
+++ b/langgraph/prebuilt/tool_executor.py
@@ -1,9 +1,11 @@
 from typing import Any, Sequence, Union
 
 from langchain_core.load.serializable import Serializable
-from langchain_core.runnables import RunnableBinding, RunnableConfig, RunnableLambda
+from langchain_core.runnables import RunnableConfig
 from langchain_core.tools import BaseTool
 
+from langgraph.utils import RunnableCallable
+
 INVALID_TOOL_MSG_TEMPLATE = (
     "{requested_tool_name} is not a valid tool, "
     "try one of [{available_tool_names_str}]."
@@ -26,29 +28,20 @@ class ToolInvocation(Serializable):
     """The input to pass in to the Tool."""
 
 
-class ToolExecutor(RunnableBinding):
-    tools: Sequence[BaseTool]
-    tool_map: dict
-    invalid_tool_msg_template: str
-
+class ToolExecutor(RunnableCallable):
     def __init__(
         self,
         tools: Sequence[BaseTool],
         *,
         invalid_tool_msg_template: str = INVALID_TOOL_MSG_TEMPLATE,
-        **kwargs: Any,
     ) -> None:
-        bound = RunnableLambda(self._execute, afunc=self._aexecute)
-        super().__init__(
-            bound=bound,
-            tools=tools,
-            tool_map={t.name: t for t in tools},
-            invalid_tool_msg_template=invalid_tool_msg_template,
-            **kwargs,
-        )
+        super().__init__(self._execute, afunc=self._aexecute, trace=False)
+        self.tools = tools
+        self.tool_map = {t.name: t for t in tools}
+        self.invalid_tool_msg_template = invalid_tool_msg_template
 
     def _execute(
-        self, tool_invocation: ToolInvocationInterface, *, config: RunnableConfig
+        self, tool_invocation: ToolInvocationInterface, config: RunnableConfig
     ) -> Any:
         if tool_invocation.tool not in self.tool_map:
             return self.invalid_tool_msg_template.format(
@@ -57,11 +50,11 @@ class ToolExecutor(RunnableBinding):
             )
         else:
             tool = self.tool_map[tool_invocation.tool]
-            output = tool.invoke(tool_invocation.tool_input, config=config)
+            output = tool.invoke(tool_invocation.tool_input, config)
             return output
 
     async def _aexecute(
-        self, tool_invocation: ToolInvocationInterface, *, config: RunnableConfig
+        self, tool_invocation: ToolInvocationInterface, config: RunnableConfig
     ) -> Any:
         if tool_invocation.tool not in self.tool_map:
             return self.invalid_tool_msg_template.format(
@@ -70,5 +63,5 @@ class ToolExecutor(RunnableBinding):
             )
         else:
             tool = self.tool_map[tool_invocation.tool]
-            output = await tool.ainvoke(tool_invocation.tool_input, config=config)
+            output = await tool.ainvoke(tool_invocation.tool_input, config)
             return output
diff --git a/tests/test_pregel.py b/tests/test_pregel.py
index 1f4703b7e..903d51dc2 100644
--- a/tests/test_pregel.py
+++ b/tests/test_pregel.py
@@ -2307,7 +2307,7 @@ def test_message_graph(
         FunctionMessage(
             content="result for query",
             name="search_api",
-            id="00000000-0000-4000-8000-000000000014",
+            id="00000000-0000-4000-8000-000000000013",
         ),
         AIMessage(
             content="",
@@ -2319,7 +2319,7 @@ def test_message_graph(
         FunctionMessage(
             content="result for another",
             name="search_api",
-            id="00000000-0000-4000-8000-000000000026",
+            id="00000000-0000-4000-8000-000000000024",
         ),
         AIMessage(content="answer", id="ai3"),
     ]
@@ -2338,7 +2338,7 @@ def test_message_graph(
             "action": FunctionMessage(
                 content="result for query",
                 name="search_api",
-                id="00000000-0000-4000-8000-000000000046",
+                id="00000000-0000-4000-8000-000000000043",
             )
         },
         {
@@ -2354,7 +2354,7 @@ def test_message_graph(
             "action": FunctionMessage(
                 content="result for another",
                 name="search_api",
-                id="00000000-0000-4000-8000-000000000058",
+                id="00000000-0000-4000-8000-000000000054",
             )
         },
         {"agent": AIMessage(content="answer", id="ai3")},