Merge branch 'nc/state' of github.com:langchain-ai/permchain into nc/state

This commit is contained in:
Harrison Chase
2024-01-15 16:11:01 -08:00
3 changed files with 74 additions and 39 deletions
+19 -18
View File
@@ -1,22 +1,17 @@
import operator
from typing import Annotated, TypedDict, Union, Sequence
from typing import Annotated, Sequence, TypedDict, Union
from langchain_core.agents import AgentAction, AgentFinish
from langchain_core.messages import BaseMessage
from langchain_core.runnables import RunnableLambda
from typing import Annotated, Optional, TypedDict
from langchain_core.agents import AgentAction, AgentFinish
from langgraph.checkpoint import BaseCheckpointSaver
from langgraph.graph import END, StateGraph
from langgraph.prebuilt.tool_executor import ToolExecutor
def _get_agent_state(input_schema= None):
def _get_agent_state(input_schema=None):
if input_schema is None:
class AgentState(TypedDict):
# The input string
input: str
@@ -31,6 +26,7 @@ def _get_agent_state(input_schema= None):
intermediate_steps: Annotated[list[tuple[AgentAction, str]], operator.add]
else:
class AgentState(input_schema):
# The outcome of a given call to the agent
# Needs `None` as a valid type, since this is what this will start as
@@ -43,12 +39,7 @@ def _get_agent_state(input_schema= None):
return AgentState
def create_agent_executor(
agent_runnable,
tools,
input_schema=None,
checkpointer: Optional[BaseCheckpointSaver] = None,
):
def create_agent_executor(agent_runnable, tools, input_schema=None):
if isinstance(tools, ToolExecutor):
tool_executor = tools
else:
@@ -73,6 +64,10 @@ def create_agent_executor(
agent_outcome = agent_runnable.invoke(data)
return {"agent_outcome": agent_outcome}
async def arun_agent(data):
agent_outcome = await agent_runnable.ainvoke(data)
return {"agent_outcome": agent_outcome}
# Define the function to execute tools
def execute_tools(data):
# Get the most recent agent_outcome - this is the key added in the `agent` above
@@ -80,12 +75,18 @@ def create_agent_executor(
output = tool_executor.invoke(agent_action)
return {"intermediate_steps": [(agent_action, str(output))]}
async def aexecute_tools(data):
# Get the most recent agent_outcome - this is the key added in the `agent` above
agent_action = data["agent_outcome"]
output = await tool_executor.ainvoke(agent_action)
return {"intermediate_steps": [(agent_action, str(output))]}
# Define a new graph
workflow = StateGraph(state)
# Define the two nodes we will cycle between
workflow.add_node("agent", run_agent)
workflow.add_node("action", execute_tools)
workflow.add_node("agent", RunnableLambda(run_agent, arun_agent))
workflow.add_node("action", RunnableLambda(execute_tools, aexecute_tools))
# Set the entrypoint as `agent`
# This means that this node is the first one called
@@ -119,4 +120,4 @@ def create_agent_executor(
# 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(checkpointer=checkpointer)
return workflow.compile()
+35 -12
View File
@@ -5,6 +5,7 @@ from typing import Annotated, Sequence, TypedDict
from langchain.tools.render import format_tool_to_openai_function
from langchain_core.agents import AgentAction
from langchain_core.messages import BaseMessage, FunctionMessage
from langchain_core.runnables import RunnableLambda
from langgraph.graph import END, StateGraph
from langgraph.prebuilt.tool_executor import ToolExecutor
@@ -17,11 +18,13 @@ def create_function_calling_executor(model, tools):
else:
tool_executor = ToolExecutor(tools)
tool_classes = tools
model = model.bind_functions([format_tool_to_openai_function(t) for t in tool_classes])
model = model.bind_functions(
[format_tool_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']
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:
@@ -32,23 +35,34 @@ def create_function_calling_executor(model, tools):
# Define the function that calls the model
def call_model(state):
messages = state['messages']
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 call_tool(state):
messages = state['messages']
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
action = AgentAction(
return AgentAction(
tool=last_message.additional_kwargs["function_call"]["name"],
tool_input=json.loads(last_message.additional_kwargs["function_call"]["arguments"]),
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
@@ -56,6 +70,15 @@ 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):
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
@@ -67,8 +90,8 @@ def create_function_calling_executor(model, tools):
workflow = StateGraph(AgentState)
# Define the two nodes we will cycle between
workflow.add_node("agent", call_model)
workflow.add_node("action", call_tool)
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
@@ -91,13 +114,13 @@ def create_function_calling_executor(model, tools):
# If `tools`, then we call the tool node.
"continue": "action",
# Otherwise we finish.
"end": END
}
"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')
workflow.add_edge("action", "agent")
# Finally, we compile it!
# This compiles it into a LangChain Runnable,
+20 -9
View File
@@ -1,9 +1,8 @@
from typing import Any, Sequence
from typing import Any, Sequence, Union
from typing import Union
from langchain_core.load.serializable import Serializable
from langchain_core.runnables import RunnableBinding, RunnableLambda
from langchain_core.tools import BaseTool
from langchain_core.load.serializable import Serializable
INVALID_TOOL_MSG_TEMPLATE = (
"{requested_tool_name} is not a valid tool, "
@@ -13,6 +12,7 @@ INVALID_TOOL_MSG_TEMPLATE = (
class ToolInvocationInterface:
"""Interface for invoking a tool"""
tool: str
tool_input: Union[str, dict]
@@ -25,21 +25,32 @@ class ToolInvocation(Serializable):
tool_input: Union[str, dict]
"""The input to pass in to the Tool."""
class ToolExecutor(RunnableBinding):
class ToolExecutor(RunnableBinding):
tools: Sequence[BaseTool]
tool_map: dict
invalid_tool_msg_template: str
def __init__(self, tools: Sequence[BaseTool], invalid_tool_msg_template: str = INVALID_TOOL_MSG_TEMPLATE, **kwargs: Any) -> None:
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__(
bound=bound,
tools=tools,
tool_map={t.name: t for t in tools},
invalid_tool_msg_template=invalid_tool_msg_template,
**kwargs,
)
def _execute(self, tool_invocation: ToolInvocationInterface) -> Any:
if tool_invocation.tool not in self.tool_map:
return self.invalid_tool_msg_template.format(
requested_tool_name=tool_invocation.tool,
available_tool_names_str=", ".join([t.name for t in self.tools])
available_tool_names_str=", ".join([t.name for t in self.tools]),
)
else:
tool = self.tool_map[tool_invocation.tool]
@@ -50,9 +61,9 @@ class ToolExecutor(RunnableBinding):
if tool_invocation.tool not in self.tool_map:
return self.invalid_tool_msg_template.format(
requested_tool_name=tool_invocation.tool,
available_tool_names_str=", ".join([t.name for t in self.tools])
available_tool_names_str=", ".join([t.name for t in self.tools]),
)
else:
tool = self.tool_map[tool_invocation.tool]
output = await tool.ainvoke(tool_invocation.tool_input)
return output
return output