Support multiple tool calls, Lint

This commit is contained in:
Nuno Campos
2024-02-19 09:12:08 -08:00
parent 6905242665
commit e4bfa8603d
3 changed files with 328 additions and 623 deletions
+70 -50
View File
@@ -1,17 +1,23 @@
import json
import operator
from typing import Annotated, Sequence, TypedDict
from typing import Annotated, Sequence, TypedDict, Union
from langchain_core.agents import AgentAction
from langchain_core.language_models import LanguageModelLike
from langchain_core.messages import BaseMessage, FunctionMessage, ToolMessage
from langchain_core.runnables import RunnableLambda
from langchain_core.utils.function_calling import convert_to_openai_function, convert_to_openai_tool
from langchain_core.tools import BaseTool
from langchain_core.utils.function_calling import (
convert_to_openai_function,
convert_to_openai_tool,
)
from langgraph.graph import END, StateGraph
from langgraph.prebuilt.tool_executor import ToolExecutor
from langgraph.prebuilt.tool_executor import ToolExecutor, ToolInvocation
def create_function_calling_executor(model, tools):
def create_function_calling_executor(
model: LanguageModelLike, tools: Union[ToolExecutor, Sequence[BaseTool]]
):
if isinstance(tools, ToolExecutor):
tool_executor = tools
tool_classes = tools.tools
@@ -20,8 +26,15 @@ def create_function_calling_executor(model, tools):
tool_classes = tools
model = model.bind(functions=[convert_to_openai_function(t) for t in tool_classes])
# We create the AgentState that we will pass around
# This simply involves a list of messages
# We want steps to return messages to append to the list
# So we annotate the messages attribute with operator.add
class AgentState(TypedDict):
messages: Annotated[Sequence[BaseMessage], operator.add]
# Define the function that determines whether to continue or not
def should_continue(state):
def should_continue(state: AgentState):
messages = state["messages"]
last_message = messages[-1]
# If there is no function call, then we finish
@@ -32,34 +45,33 @@ def create_function_calling_executor(model, tools):
return "continue"
# Define the function that calls the model
def call_model(state):
def call_model(state: AgentState):
messages = state["messages"]
response = model.invoke(messages)
# We return a list, because this will get added to the existing list
return {"messages": [response]}
async def acall_model(state):
async def acall_model(state: AgentState):
messages = state["messages"]
response = await model.ainvoke(messages)
# We return a list, because this will get added to the existing list
return {"messages": [response]}
# Define the function to execute tools
def _get_action(state):
def _get_action(state: AgentState):
messages = state["messages"]
# Based on the continue condition
# we know the last message involves a function call
last_message = messages[-1]
# We construct an AgentAction from the function_call
return AgentAction(
return ToolInvocation(
tool=last_message.additional_kwargs["function_call"]["name"],
tool_input=json.loads(
last_message.additional_kwargs["function_call"]["arguments"]
),
log="",
)
def call_tool(state):
def call_tool(state: AgentState):
action = _get_action(state)
# We call the tool_executor and get back a response
response = tool_executor.invoke(action)
@@ -68,7 +80,7 @@ def create_function_calling_executor(model, tools):
# We return a list, because this will get added to the existing list
return {"messages": [function_message]}
async def acall_tool(state):
async def acall_tool(state: AgentState):
action = _get_action(state)
# We call the tool_executor and get back a response
response = await tool_executor.ainvoke(action)
@@ -77,13 +89,6 @@ def create_function_calling_executor(model, tools):
# We return a list, because this will get added to the existing list
return {"messages": [function_message]}
# We create the AgentState that we will pass around
# This simply involves a list of messages
# We want steps to return messages to append to the list
# So we annotate the messages attribute with operator.add
class AgentState(TypedDict):
messages: Annotated[Sequence[BaseMessage], operator.add]
# Define a new graph
workflow = StateGraph(AgentState)
@@ -125,17 +130,27 @@ def create_function_calling_executor(model, tools):
# meaning you can use it as you would any other runnable
return workflow.compile()
def create_tool_calling_executor(model, tools):
def create_tool_calling_executor(
model: LanguageModelLike, tools: Union[ToolExecutor, Sequence[BaseTool]]
):
if isinstance(tools, ToolExecutor):
tool_executor = tools
tool_classes = tools.tools
else:
tool_executor = ToolExecutor(tools)
tool_classes = tools
model = model.bind(functions=[convert_to_openai_tool(t) for t in tool_classes])
model = model.bind(tools=[convert_to_openai_tool(t) for t in tool_classes])
# We create the AgentState that we will pass around
# This simply involves a list of messages
# We want steps to return messages to append to the list
# So we annotate the messages attribute with operator.add
class AgentState(TypedDict):
messages: Annotated[Sequence[BaseMessage], operator.add]
# Define the function that determines whether to continue or not
def should_continue(state):
def should_continue(state: AgentState):
messages = state["messages"]
last_message = messages[-1]
# If there is no function call, then we finish
@@ -146,57 +161,62 @@ def create_tool_calling_executor(model, tools):
return "continue"
# Define the function that calls the model
def call_model(state):
def call_model(state: AgentState):
messages = state["messages"]
response = model.invoke(messages)
# We return a list, because this will get added to the existing list
return {"messages": [response]}
async def acall_model(state):
async def acall_model(state: AgentState):
messages = state["messages"]
response = await model.ainvoke(messages)
# We return a list, because this will get added to the existing list
return {"messages": [response]}
# Define the function to execute tools
def _get_action(state):
def _get_actions(state: AgentState):
messages = state["messages"]
# Based on the continue condition
# we know the last message involves a tool call
last_message = messages[-1]
# We construct an AgentAction from the tool_calls
return AgentAction(
tool=last_message.additional_kwargs["tool_calls"][0]["function"]["name"],
tool_input=json.loads(
last_message.additional_kwargs["tool_calls"][0]["function"]["arguments"]
),
log=last_message.additional_kwargs["tool_calls"][0]["id"],
# We construct an AgentAction from each of the tool_calls
return (
[
ToolInvocation(
tool=tool_call["function"]["name"],
tool_input=json.loads(tool_call["function"]["arguments"]),
)
for tool_call in last_message.additional_kwargs["tool_calls"]
],
[
tool_call["id"]
for tool_call in last_message.additional_kwargs["tool_calls"]
],
)
def call_tool(state):
action = _get_action(state)
def call_tool(state: AgentState):
actions, ids = _get_actions(state)
# We call the tool_executor and get back a response
response = tool_executor.invoke(action)
responses = tool_executor.batch(actions)
# We use the response to create a FunctionMessage
tool_message = ToolMessage(content=str(response), tool_call_id=action.log)
tool_messages = [
ToolMessage(content=str(response), tool_call_id=id)
for response, id in zip(responses, ids)
]
# We return a list, because this will get added to the existing list
return {"messages": [tool_message]}
return {"messages": tool_messages}
async def acall_tool(state):
action = _get_action(state)
async def acall_tool(state: AgentState):
actions, ids = _get_actions(state)
# We call the tool_executor and get back a response
response = await tool_executor.ainvoke(action)
responses = await tool_executor.abatch(actions)
# We use the response to create a FunctionMessage
tool_message = ToolMessage(content=str(response), tool_call_id=action.log)
tool_messages = [
ToolMessage(content=str(response), tool_call_id=id)
for response, id in zip(responses, ids)
]
# We return a list, because this will get added to the existing list
return {"messages": [tool_message]}
# We create the AgentState that we will pass around
# This simply involves a list of messages
# We want steps to return messages to append to the list
# So we annotate the messages attribute with operator.add
class AgentState(TypedDict):
messages: Annotated[Sequence[BaseMessage], operator.add]
return {"messages": tool_messages}
# Define a new graph
workflow = StateGraph(AgentState)
+127 -281
View File
@@ -20,11 +20,15 @@ from langgraph.checkpoint.sqlite import SqliteSaver
from langgraph.graph import END, Graph
from langgraph.graph.message import MessageGraph
from langgraph.graph.state import StateGraph
from langgraph.prebuilt.chat_agent_executor import create_function_calling_executor, create_tool_calling_executor
from langgraph.prebuilt.chat_agent_executor import (
create_function_calling_executor,
create_tool_calling_executor,
)
from langgraph.prebuilt.tool_executor import ToolExecutor
from langgraph.pregel import Channel, GraphRecursionError, Pregel
from langgraph.pregel.reserved import ReservedChannels
def test_invoke_single_process_in_out(mocker: MockerFixture) -> None:
add_one = mocker.Mock(side_effect=lambda x: x + 1)
chain = Channel.subscribe_to("input") | add_one | Channel.write_to("output")
@@ -1060,6 +1064,7 @@ def test_conditional_graph_state() -> None:
},
]
def test_prebuilt_tool_chat() -> None:
from langchain.chat_models.fake import FakeMessagesListChatModel
from langchain_community.tools import tool
@@ -1082,27 +1087,39 @@ def test_prebuilt_tool_chat() -> None:
AIMessage(
content="",
additional_kwargs={
"tool_calls": [{
"id": "tool_call123",
"type": "function",
"function":{
"name": "search_api",
"arguments": json.dumps("query"),
"tool_calls": [
{
"id": "tool_call123",
"type": "function",
"function": {
"name": "search_api",
"arguments": json.dumps("query"),
},
}
}]
]
},
),
AIMessage(
content="",
additional_kwargs={
"tool_calls": [{
"id": "tool_call234",
"type": "function",
"function":{
"name": "search_api",
"arguments": json.dumps("another"),
}
}]
"tool_calls": [
{
"id": "tool_call234",
"type": "function",
"function": {
"name": "search_api",
"arguments": json.dumps("another"),
},
},
{
"id": "tool_call567",
"type": "function",
"function": {
"name": "search_api",
"arguments": '"a third one"',
},
},
]
},
),
AIMessage(content="answer"),
@@ -1110,7 +1127,7 @@ def test_prebuilt_tool_chat() -> None:
),
tools,
)
assert app.invoke(
{"messages": [HumanMessage(content="what is weather in sf")]}
) == {
@@ -1119,31 +1136,44 @@ def test_prebuilt_tool_chat() -> None:
AIMessage(
content="",
additional_kwargs={
"tool_calls": [{
"id": "tool_call123",
"type": "function",
"function":{
"name": "search_api",
"arguments": "\"query\"",
"tool_calls": [
{
"id": "tool_call123",
"type": "function",
"function": {
"name": "search_api",
"arguments": '"query"',
},
}
}]
]
},
),
ToolMessage(content="result for query", tool_call_id="tool_call123"),
AIMessage(
content="",
additional_kwargs={
"tool_calls": [{
"id": "tool_call234",
"type": "function",
"function":{
"name": "search_api",
"arguments": "\"another\"",
}
}]
},
"tool_calls": [
{
"id": "tool_call234",
"type": "function",
"function": {
"name": "search_api",
"arguments": '"another"',
},
},
{
"id": "tool_call567",
"type": "function",
"function": {
"name": "search_api",
"arguments": '"a third one"',
},
},
]
},
),
ToolMessage(content="result for another", tool_call_id="tool_call234"),
ToolMessage(content="result for a third one", tool_call_id="tool_call567"),
AIMessage(content="answer"),
]
}
@@ -1157,14 +1187,16 @@ def test_prebuilt_tool_chat() -> None:
AIMessage(
content="",
additional_kwargs={
"tool_calls": [{
"tool_calls": [
{
"id": "tool_call123",
"type": "function",
"function":{
"function": {
"name": "search_api",
"arguments": "\"query\"",
}
}]
"arguments": '"query"',
},
}
]
},
)
]
@@ -1182,16 +1214,26 @@ def test_prebuilt_tool_chat() -> None:
"messages": [
AIMessage(
content="",
additional_kwargs={
"tool_calls": [{
additional_kwargs={
"tool_calls": [
{
"id": "tool_call234",
"type": "function",
"function":{
"function": {
"name": "search_api",
"arguments": "\"another\"",
}
}]
},
"arguments": '"another"',
},
},
{
"id": "tool_call567",
"type": "function",
"function": {
"name": "search_api",
"arguments": '"a third one"',
},
},
]
},
)
]
}
@@ -1199,7 +1241,12 @@ def test_prebuilt_tool_chat() -> None:
{
"action": {
"messages": [
ToolMessage(content="result for another", tool_call_id="tool_call234")
ToolMessage(
content="result for another", tool_call_id="tool_call234"
),
ToolMessage(
content="result for a third one", tool_call_id="tool_call567"
),
]
}
},
@@ -1211,258 +1258,56 @@ def test_prebuilt_tool_chat() -> None:
AIMessage(
content="",
additional_kwargs={
"tool_calls": [{
"tool_calls": [
{
"id": "tool_call123",
"type": "function",
"function":{
"function": {
"name": "search_api",
"arguments": "\"query\"",
}
}]
},
"arguments": '"query"',
},
}
]
},
),
ToolMessage(
content="result for query", tool_call_id="tool_call123"
),
ToolMessage(content="result for query", tool_call_id="tool_call123"),
AIMessage(
content="",
additional_kwargs={
"tool_calls": [{
"tool_calls": [
{
"id": "tool_call234",
"type": "function",
"function":{
"function": {
"name": "search_api",
"arguments": "\"another\"",
}
}]
"arguments": '"another"',
},
},
{
"id": "tool_call567",
"type": "function",
"function": {
"name": "search_api",
"arguments": '"a third one"',
},
},
]
},
),
ToolMessage(content="result for another", tool_call_id="tool_call234"),
ToolMessage(
content="result for another", tool_call_id="tool_call234"
),
ToolMessage(
content="result for a third one", tool_call_id="tool_call567"
),
AIMessage(content="answer"),
]
}
},
]
def test_tool_message_graph() -> None:
from langchain.chat_models.fake import FakeMessagesListChatModel
from langchain_community.tools import tool
from langchain_core.agents import AgentAction
from langchain_core.messages import AIMessage, ToolMessage, HumanMessage
class FakeFuntionChatModel(FakeMessagesListChatModel):
def bind_functions(self, functions: list):
return self
@tool()
def search_api(query: str) -> str:
"""Searches the API for the query."""
return f"result for {query}"
tools = [search_api]
model = FakeFuntionChatModel(
responses=[
AIMessage(
content="",
additional_kwargs={
"tool_calls": [{
"id": "tool_call123",
"type": "function",
"function":{
"name": "search_api",
"arguments": json.dumps("query"),
}
}]
},
),
AIMessage(
content="",
additional_kwargs={
"tool_calls": [{
"id": "tool_call234",
"type": "function",
"function":{
"name": "search_api",
"arguments": json.dumps("another"),
}
}]
},
),
AIMessage(content="answer"),
]
)
tool_executor = ToolExecutor(tools)
# Define the function that determines whether to continue or not
def should_continue(messages):
last_message = messages[-1]
# If there is no function call, then we finish
if "tool_calls" not in last_message.additional_kwargs:
return "end"
# Otherwise if there is, we continue
else:
return "continue"
def call_tool(messages):
# Based on the continue condition
# we know the last message involves a function call
last_message = messages[-1]
# We construct an AgentAction from the function_call
action = AgentAction(
tool=last_message.additional_kwargs["tool_calls"][0]["function"]["name"],
tool_input=json.loads(
last_message.additional_kwargs["tool_calls"][0]["function"]["arguments"]
),
log=last_message.additional_kwargs["tool_calls"][0]["id"],
)
# We call the tool_executor and get back a response
response = tool_executor.invoke(action)
# We use the response to create a ToolMessage
return ToolMessage(content=str(response), tool_call_id=action.log)
# Define a new graph
workflow = MessageGraph()
# Define the two nodes we will cycle between
workflow.add_node("agent", model)
workflow.add_node("action", call_tool)
# Set the entrypoint as `agent`
# This means that this node is the first one called
workflow.set_entry_point("agent")
# We now add a conditional edge
workflow.add_conditional_edges(
# First, we define the start node. We use `agent`.
# This means these are the edges taken after the `agent` node is called.
"agent",
# Next, we pass in the function that will determine which node is called next.
should_continue,
# Finally we pass in a mapping.
# The keys are strings, and the values are other nodes.
# END is a special node marking that the graph should finish.
# What will happen is we will call `should_continue`, and then the output of that
# will be matched against the keys in this mapping.
# Based on which one it matches, that node will then be called.
{
# If `tools`, then we call the tool node.
"continue": "action",
# Otherwise we finish.
"end": END,
},
)
# We now add a normal edge from `tools` to `agent`.
# This means that after `tools` is called, `agent` node is called next.
workflow.add_edge("action", "agent")
# Finally, we compile it!
# This compiles it into a LangChain Runnable,
# meaning you can use it as you would any other runnable
app = workflow.compile()
assert app.invoke(HumanMessage(content="what is weather in sf")) == [
HumanMessage(content="what is weather in sf"),
AIMessage(
content="",
additional_kwargs={
"tool_calls": [{
"id": "tool_call123",
"type": "function",
"function":{
"name": "search_api",
"arguments": "\"query\"",
}
}]
},
),
ToolMessage(content="result for query", tool_call_id="tool_call123"),
AIMessage(
content="",
additional_kwargs={
"tool_calls": [{
"id": "tool_call234",
"type": "function",
"function":{
"name": "search_api",
"arguments": "\"another\"",
}
}]
},
),
ToolMessage(content="result for another", tool_call_id="tool_call234"),
AIMessage(content="answer"),
]
assert [*app.stream([HumanMessage(content="what is weather in sf")])] == [
{
"agent": AIMessage(
content="",
additional_kwargs={
"tool_calls": [{
"id": "tool_call123",
"type": "function",
"function":{
"name": "search_api",
"arguments": "\"query\"",
}
}]
},
)
},
{"action": ToolMessage(content="result for query", tool_call_id="tool_call123")},
{
"agent": AIMessage(
content="",
additional_kwargs={
"tool_calls": [{
"id": "tool_call234",
"type": "function",
"function":{
"name": "search_api",
"arguments": "\"another\"",
}
}]
},
)
},
{"action": ToolMessage(content="result for another", tool_call_id="tool_call234")},
{"agent": AIMessage(content="answer")},
{
"__end__": [
HumanMessage(content="what is weather in sf"),
AIMessage(
content="",
additional_kwargs={
"tool_calls": [{
"id": "tool_call123",
"type": "function",
"function":{
"name": "search_api",
"arguments": "\"query\"",
}
}]
},
),
ToolMessage(content="result for query", tool_call_id="tool_call123"),
AIMessage(
content="",
additional_kwargs={
"tool_calls": [{
"id": "tool_call234",
"type": "function",
"function":{
"name": "search_api",
"arguments": "\"another\"",
}
}]
},
),
ToolMessage(content="result for another", tool_call_id="tool_call234"),
AIMessage(content="answer"),
]
},
]
def test_prebuilt_chat() -> None:
from langchain.chat_models.fake import FakeMessagesListChatModel
@@ -1608,6 +1453,7 @@ def test_prebuilt_chat() -> None:
},
]
def test_message_graph() -> None:
from langchain.chat_models.fake import FakeMessagesListChatModel
from langchain_community.tools import tool
@@ -1781,4 +1627,4 @@ def test_message_graph() -> None:
AIMessage(content="answer"),
]
},
]
]
+131 -292
View File
@@ -26,7 +26,10 @@ from langgraph.checkpoint.aiosqlite import AsyncSqliteSaver
from langgraph.checkpoint.memory import MemorySaver
from langgraph.graph import END, Graph, StateGraph
from langgraph.graph.message import MessageGraph
from langgraph.prebuilt.chat_agent_executor import create_function_calling_executor, create_tool_calling_executor
from langgraph.prebuilt.chat_agent_executor import (
create_function_calling_executor,
create_tool_calling_executor,
)
from langgraph.prebuilt.tool_executor import ToolExecutor
from langgraph.pregel import Channel, GraphRecursionError, Pregel
from langgraph.pregel.reserved import ReservedChannels
@@ -1114,7 +1117,7 @@ async def test_conditional_graph_state() -> None:
async def test_prebuilt_tool_chat() -> None:
from langchain.chat_models.fake import FakeMessagesListChatModel
from langchain_community.tools import tool
from langchain_core.messages import AIMessage, ToolMessage, HumanMessage
from langchain_core.messages import AIMessage, HumanMessage, ToolMessage
class FakeFuntionChatModel(FakeMessagesListChatModel):
def bind_functions(self, functions: list):
@@ -1133,27 +1136,39 @@ async def test_prebuilt_tool_chat() -> None:
AIMessage(
content="",
additional_kwargs={
"tool_calls": [{
"id": "tool_call123",
"type": "function",
"function":{
"name": "search_api",
"arguments": json.dumps("query"),
"tool_calls": [
{
"id": "tool_call123",
"type": "function",
"function": {
"name": "search_api",
"arguments": json.dumps("query"),
},
}
}]
]
},
),
AIMessage(
content="",
additional_kwargs={
"tool_calls": [{
"id": "tool_call234",
"type": "function",
"function":{
"name": "search_api",
"arguments": json.dumps("another"),
}
}]
"tool_calls": [
{
"id": "tool_call234",
"type": "function",
"function": {
"name": "search_api",
"arguments": json.dumps("another"),
},
},
{
"id": "tool_call567",
"type": "function",
"function": {
"name": "search_api",
"arguments": '"a third one"',
},
},
]
},
),
AIMessage(content="answer"),
@@ -1170,31 +1185,44 @@ async def test_prebuilt_tool_chat() -> None:
AIMessage(
content="",
additional_kwargs={
"tool_calls": [{
"id": "tool_call123",
"type": "function",
"function":{
"name": "search_api",
"arguments": "\"query\"",
"tool_calls": [
{
"id": "tool_call123",
"type": "function",
"function": {
"name": "search_api",
"arguments": '"query"',
},
}
}]
},
]
},
),
ToolMessage(content="result for query", tool_call_id="tool_call123"),
AIMessage(
content="",
additional_kwargs={
"tool_calls": [{
"id": "tool_call234",
"type": "function",
"function":{
"name": "search_api",
"arguments": "\"another\"",
}
}]
"tool_calls": [
{
"id": "tool_call234",
"type": "function",
"function": {
"name": "search_api",
"arguments": '"another"',
},
},
{
"id": "tool_call567",
"type": "function",
"function": {
"name": "search_api",
"arguments": '"a third one"',
},
},
]
},
),
ToolMessage(content="result for another", tool_call_id="tool_call234"),
ToolMessage(content="result for a third one", tool_call_id="tool_call567"),
AIMessage(content="answer"),
]
}
@@ -1211,15 +1239,17 @@ async def test_prebuilt_tool_chat() -> None:
AIMessage(
content="",
additional_kwargs={
"tool_calls": [{
"id": "tool_call123",
"type": "function",
"function":{
"name": "search_api",
"arguments": "\"query\"",
"tool_calls": [
{
"id": "tool_call123",
"type": "function",
"function": {
"name": "search_api",
"arguments": '"query"',
},
}
}]
},
]
},
)
]
}
@@ -1237,14 +1267,24 @@ async def test_prebuilt_tool_chat() -> None:
AIMessage(
content="",
additional_kwargs={
"tool_calls": [{
"id": "tool_call234",
"type": "function",
"function":{
"name": "search_api",
"arguments": "\"another\"",
}
}]
"tool_calls": [
{
"id": "tool_call234",
"type": "function",
"function": {
"name": "search_api",
"arguments": '"another"',
},
},
{
"id": "tool_call567",
"type": "function",
"function": {
"name": "search_api",
"arguments": '"a third one"',
},
},
]
},
)
]
@@ -1253,7 +1293,12 @@ async def test_prebuilt_tool_chat() -> None:
{
"action": {
"messages": [
ToolMessage(content="result for another", tool_call_id="tool_call234")
ToolMessage(
content="result for another", tool_call_id="tool_call234"
),
ToolMessage(
content="result for a third one", tool_call_id="tool_call567"
),
]
}
},
@@ -1265,31 +1310,50 @@ async def test_prebuilt_tool_chat() -> None:
AIMessage(
content="",
additional_kwargs={
"tool_calls": [{
"id": "tool_call123",
"type": "function",
"function":{
"name": "search_api",
"arguments": "\"query\"",
"tool_calls": [
{
"id": "tool_call123",
"type": "function",
"function": {
"name": "search_api",
"arguments": '"query"',
},
}
}]
]
},
),
ToolMessage(content="result for query", tool_call_id="tool_call123"),
ToolMessage(
content="result for query", tool_call_id="tool_call123"
),
AIMessage(
content="",
additional_kwargs={
"tool_calls": [{
"id": "tool_call234",
"type": "function",
"function":{
"name": "search_api",
"arguments": "\"another\"",
}
}]
"tool_calls": [
{
"id": "tool_call234",
"type": "function",
"function": {
"name": "search_api",
"arguments": '"another"',
},
},
{
"id": "tool_call567",
"type": "function",
"function": {
"name": "search_api",
"arguments": '"a third one"',
},
},
]
},
),
ToolMessage(content="result for another", tool_call_id="tool_call234"),
ToolMessage(
content="result for another", tool_call_id="tool_call234"
),
ToolMessage(
content="result for a third one", tool_call_id="tool_call567"
),
AIMessage(content="answer"),
]
}
@@ -1297,231 +1361,6 @@ async def test_prebuilt_tool_chat() -> None:
]
async def test_message_tool_graph() -> None:
from langchain.chat_models.fake import FakeMessagesListChatModel
from langchain_community.tools import tool
from langchain_core.agents import AgentAction
from langchain_core.messages import AIMessage, ToolMessage, HumanMessage
class FakeFuntionChatModel(FakeMessagesListChatModel):
def bind_functions(self, functions: list):
return self
@tool()
def search_api(query: str) -> str:
"""Searches the API for the query."""
return f"result for {query}"
tools = [search_api]
model = FakeFuntionChatModel(
responses=[
AIMessage(
content="",
additional_kwargs={
"tool_calls": [{
"id": "tool_call123",
"type": "function",
"function":{
"name": "search_api",
"arguments": json.dumps("query"),
}
}]
},
),
AIMessage(
content="",
additional_kwargs={
"tool_calls": [{
"id": "tool_call234",
"type": "function",
"function":{
"name": "search_api",
"arguments": json.dumps("another"),
}
}]
},
),
AIMessage(content="answer"),
]
)
tool_executor = ToolExecutor(tools)
# Define the function that determines whether to continue or not
def should_continue(messages):
last_message = messages[-1]
# If there is no function call, then we finish
if "tool_calls" not in last_message.additional_kwargs:
return "end"
# Otherwise if there is, we continue
else:
return "continue"
async def call_tool(messages):
# Based on the continue condition
# we know the last message involves a function call
last_message = messages[-1]
# We construct an AgentAction from the function_call
action = AgentAction(
tool=last_message.additional_kwargs["tool_calls"][0]["function"]["name"],
tool_input=json.loads(
last_message.additional_kwargs["tool_calls"][0]["function"]["arguments"]
),
log=last_message.additional_kwargs["tool_calls"][0]["id"],
)
# We call the tool_executor and get back a response
response = await tool_executor.ainvoke(action)
# We use the response to create a FunctionMessage
return ToolMessage(content=str(response), tool_call_id=action.log)
# Define a new graph
workflow = MessageGraph()
# Define the two nodes we will cycle between
workflow.add_node("agent", model)
workflow.add_node("action", call_tool)
# Set the entrypoint as `agent`
# This means that this node is the first one called
workflow.set_entry_point("agent")
# We now add a conditional edge
workflow.add_conditional_edges(
# First, we define the start node. We use `agent`.
# This means these are the edges taken after the `agent` node is called.
"agent",
# Next, we pass in the function that will determine which node is called next.
should_continue,
# Finally we pass in a mapping.
# The keys are strings, and the values are other nodes.
# END is a special node marking that the graph should finish.
# What will happen is we will call `should_continue`, and then the output of that
# will be matched against the keys in this mapping.
# Based on which one it matches, that node will then be called.
{
# If `tools`, then we call the tool node.
"continue": "action",
# Otherwise we finish.
"end": END,
},
)
# We now add a normal edge from `tools` to `agent`.
# This means that after `tools` is called, `agent` node is called next.
workflow.add_edge("action", "agent")
# Finally, we compile it!
# This compiles it into a LangChain Runnable,
# meaning you can use it as you would any other runnable
app = workflow.compile()
assert await app.ainvoke(HumanMessage(content="what is weather in sf")) == [
HumanMessage(content="what is weather in sf"),
AIMessage(
content="",
additional_kwargs={
"tool_calls": [{
"id": "tool_call123",
"type": "function",
"function":{
"name": "search_api",
"arguments": "\"query\"",
}
}]
},
),
ToolMessage(content="result for query", tool_call_id="tool_call123"),
AIMessage(
content="",
additional_kwargs={
"tool_calls": [{
"id": "tool_call234",
"type": "function",
"function":{
"name": "search_api",
"arguments": "\"another\"",
}
}]
},
),
ToolMessage(content="result for another", tool_call_id="tool_call234"),
AIMessage(content="answer"),
]
assert [
c async for c in app.astream([HumanMessage(content="what is weather in sf")])
] == [
{
"agent": AIMessage(
content="",
additional_kwargs={
"tool_calls": [{
"id": "tool_call123",
"type": "function",
"function":{
"name": "search_api",
"arguments": "\"query\"",
}
}]
},
)
},
{"action": ToolMessage(content="result for query", tool_call_id="tool_call123")},
{
"agent": AIMessage(
content="",
additional_kwargs={
"tool_calls": [{
"id": "tool_call234",
"type": "function",
"function":{
"name": "search_api",
"arguments": "\"another\"",
}
}]
},
)
},
{"action": ToolMessage(content="result for another", tool_call_id="tool_call234")},
{"agent": AIMessage(content="answer")},
{
"__end__": [
HumanMessage(content="what is weather in sf"),
AIMessage(
content="",
additional_kwargs={
"tool_calls": [{
"id": "tool_call123",
"type": "function",
"function":{
"name": "search_api",
"arguments": "\"query\"",
}
}]
},
),
ToolMessage(content="result for query", tool_call_id="tool_call123"),
AIMessage(
content="",
additional_kwargs={
"tool_calls": [{
"id": "tool_call234",
"type": "function",
"function":{
"name": "search_api",
"arguments": "\"another\"",
}
}]
},
),
ToolMessage(content="result for another", tool_call_id="tool_call234"),
AIMessage(content="answer"),
]
},
]
async def test_prebuilt_chat() -> None:
from langchain.chat_models.fake import FakeMessagesListChatModel
from langchain_community.tools import tool