Simplify tool executor

This commit is contained in:
Nuno Campos
2024-04-02 14:46:14 -07:00
parent e3a9760bb2
commit 2f8ae4fdb4
2 changed files with 16 additions and 23 deletions
+12 -19
View File
@@ -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
+4 -4
View File
@@ -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")},