mirror of
https://github.com/langchain-ai/langgraph.git
synced 2026-08-20 06:35:46 +02:00
Simplify tool executor
This commit is contained in:
@@ -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
|
||||
|
||||
@@ -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")},
|
||||
|
||||
Reference in New Issue
Block a user