mirror of
https://github.com/langchain-ai/langgraph.git
synced 2026-08-30 11:49:38 +02:00
127 lines
5.0 KiB
Python
127 lines
5.0 KiB
Python
import json
|
|
import operator
|
|
from typing import Annotated, Sequence, TypedDict
|
|
|
|
from langchain_core.agents import AgentAction
|
|
from langchain_core.messages import BaseMessage, FunctionMessage
|
|
from langchain_core.runnables import RunnableLambda
|
|
from langchain_core.utils.function_calling import convert_to_openai_function
|
|
|
|
from langgraph.graph import END, StateGraph
|
|
from langgraph.prebuilt.tool_executor import ToolExecutor
|
|
|
|
|
|
def create_function_calling_executor(model, tools):
|
|
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_function(t) for t in tool_classes])
|
|
|
|
# Define the function that determines whether to continue or not
|
|
def should_continue(state):
|
|
messages = state["messages"]
|
|
last_message = messages[-1]
|
|
# If there is no function call, then we finish
|
|
if "function_call" not in last_message.additional_kwargs:
|
|
return "end"
|
|
# Otherwise if there is, we continue
|
|
else:
|
|
return "continue"
|
|
|
|
# Define the function that calls the model
|
|
def call_model(state):
|
|
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):
|
|
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):
|
|
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(
|
|
tool=last_message.additional_kwargs["function_call"]["name"],
|
|
tool_input=json.loads(
|
|
last_message.additional_kwargs["function_call"]["arguments"]
|
|
),
|
|
log="",
|
|
)
|
|
|
|
def call_tool(state):
|
|
action = _get_action(state)
|
|
# We call the tool_executor and get back a response
|
|
response = tool_executor.invoke(action)
|
|
# We use the response to create a FunctionMessage
|
|
function_message = FunctionMessage(content=str(response), name=action.tool)
|
|
# We return a list, because this will get added to the existing list
|
|
return {"messages": [function_message]}
|
|
|
|
async def acall_tool(state):
|
|
action = _get_action(state)
|
|
# We call the tool_executor and get back a response
|
|
response = await tool_executor.ainvoke(action)
|
|
# We use the response to create a FunctionMessage
|
|
function_message = FunctionMessage(content=str(response), name=action.tool)
|
|
# 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)
|
|
|
|
# Define the two nodes we will cycle between
|
|
workflow.add_node("agent", RunnableLambda(call_model, acall_model))
|
|
workflow.add_node("action", RunnableLambda(call_tool, acall_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
|
|
return workflow.compile()
|