From d12c2bae6b4a9aeebc3953fa65e9b312504f094f Mon Sep 17 00:00:00 2001 From: Harrison Chase Date: Mon, 29 Apr 2024 14:58:25 -0700 Subject: [PATCH 1/4] add arguments to chat agent executor --- langgraph/graph/graph.py | 4 +--- langgraph/prebuilt/chat_agent_executor.py | 17 ++++++++++++++--- langgraph/pregel/__init__.py | 12 ++++-------- langgraph/serde/base.py | 6 ++---- langgraph/version.py | 1 + 5 files changed, 22 insertions(+), 18 deletions(-) diff --git a/langgraph/graph/graph.py b/langgraph/graph/graph.py index 7291b4e15..36cc01cc0 100644 --- a/langgraph/graph/graph.py +++ b/langgraph/graph/graph.py @@ -15,9 +15,7 @@ from typing import ( from langchain_core.runnables import Runnable from langchain_core.runnables.base import RunnableLike from langchain_core.runnables.config import RunnableConfig -from langchain_core.runnables.graph import ( - Node as RunnableGraphNode, -) +from langchain_core.runnables.graph import Node as RunnableGraphNode from langgraph.channels.ephemeral_value import EphemeralValue from langgraph.checkpoint import BaseCheckpointSaver diff --git a/langgraph/prebuilt/chat_agent_executor.py b/langgraph/prebuilt/chat_agent_executor.py index fbc254b96..03d90c890 100644 --- a/langgraph/prebuilt/chat_agent_executor.py +++ b/langgraph/prebuilt/chat_agent_executor.py @@ -1,5 +1,5 @@ import json -from typing import Annotated, Sequence, TypedDict, Union +from typing import Annotated, Optional, Sequence, TypedDict, Union from langchain_core.language_models import LanguageModelLike from langchain_core.messages import BaseMessage, FunctionMessage @@ -7,6 +7,7 @@ from langchain_core.runnables import RunnableLambda from langchain_core.tools import BaseTool from langchain_core.utils.function_calling import convert_to_openai_function +from langgraph.checkpoint import BaseCheckpointSaver from langgraph.graph import END, StateGraph from langgraph.graph.graph import CompiledGraph from langgraph.graph.message import add_messages @@ -134,7 +135,12 @@ def create_function_calling_executor( def create_tool_calling_executor( - model: LanguageModelLike, tools: Union[ToolExecutor, Sequence[BaseTool]] + model: LanguageModelLike, + tools: Union[ToolExecutor, Sequence[BaseTool]], + checkpointer: Optional[BaseCheckpointSaver] = None, + interrupt_before: Optional[Sequence[str]] = None, + interrupt_after: Optional[Sequence[str]] = None, + debug: bool = False, ) -> CompiledGraph: """Creates a graph that works with a chat model that utilizes tool calling. @@ -231,4 +237,9 @@ def create_tool_calling_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() + return workflow.compile( + checkpointer=checkpointer, + interrupt_before=interrupt_before, + interrupt_after=interrupt_after, + debug=debug, + ) diff --git a/langgraph/pregel/__init__.py b/langgraph/pregel/__init__.py index 5d5f90633..c26b05241 100644 --- a/langgraph/pregel/__init__.py +++ b/langgraph/pregel/__init__.py @@ -116,8 +116,7 @@ class Channel: *, key: Optional[str] = None, tags: Optional[list[str]] = None, - ) -> PregelNode: - ... + ) -> PregelNode: ... @overload @classmethod @@ -127,8 +126,7 @@ class Channel: *, key: None = None, tags: Optional[list[str]] = None, - ) -> PregelNode: - ... + ) -> PregelNode: ... @classmethod def subscribe_to( @@ -1212,8 +1210,7 @@ def _prepare_next_tasks( processes: Mapping[str, PregelNode], channels: Mapping[str, BaseChannel], for_execution: Literal[False], -) -> tuple[Checkpoint, list[PregelTaskDescription]]: - ... +) -> tuple[Checkpoint, list[PregelTaskDescription]]: ... @overload @@ -1222,8 +1219,7 @@ def _prepare_next_tasks( processes: Mapping[str, PregelNode], channels: Mapping[str, BaseChannel], for_execution: Literal[True], -) -> tuple[Checkpoint, list[PregelExecutableTask]]: - ... +) -> tuple[Checkpoint, list[PregelExecutableTask]]: ... def _prepare_next_tasks( diff --git a/langgraph/serde/base.py b/langgraph/serde/base.py index 2fbb0ab71..5f1250d1e 100644 --- a/langgraph/serde/base.py +++ b/langgraph/serde/base.py @@ -10,8 +10,6 @@ class SerializerProtocol(Protocol): Valid implementations include the `pickle`, `json` and `orjson` modules. """ - def dumps(self, obj: Any) -> bytes: - ... + def dumps(self, obj: Any) -> bytes: ... - def loads(self, data: bytes) -> Any: - ... + def loads(self, data: bytes) -> Any: ... diff --git a/langgraph/version.py b/langgraph/version.py index ac7aeef6f..3368893c0 100644 --- a/langgraph/version.py +++ b/langgraph/version.py @@ -1,4 +1,5 @@ """Main entrypoint into package.""" + from importlib import metadata try: From a32fe442f0eb9ceff5d70a3f4412a98cde512900 Mon Sep 17 00:00:00 2001 From: Harrison Chase Date: Mon, 29 Apr 2024 14:59:31 -0700 Subject: [PATCH 2/4] cr --- langgraph/graph/graph.py | 4 +++- langgraph/pregel/__init__.py | 12 ++++++++---- langgraph/serde/base.py | 6 ++++-- langgraph/version.py | 1 - 4 files changed, 15 insertions(+), 8 deletions(-) diff --git a/langgraph/graph/graph.py b/langgraph/graph/graph.py index 36cc01cc0..7291b4e15 100644 --- a/langgraph/graph/graph.py +++ b/langgraph/graph/graph.py @@ -15,7 +15,9 @@ from typing import ( from langchain_core.runnables import Runnable from langchain_core.runnables.base import RunnableLike from langchain_core.runnables.config import RunnableConfig -from langchain_core.runnables.graph import Node as RunnableGraphNode +from langchain_core.runnables.graph import ( + Node as RunnableGraphNode, +) from langgraph.channels.ephemeral_value import EphemeralValue from langgraph.checkpoint import BaseCheckpointSaver diff --git a/langgraph/pregel/__init__.py b/langgraph/pregel/__init__.py index c26b05241..5d5f90633 100644 --- a/langgraph/pregel/__init__.py +++ b/langgraph/pregel/__init__.py @@ -116,7 +116,8 @@ class Channel: *, key: Optional[str] = None, tags: Optional[list[str]] = None, - ) -> PregelNode: ... + ) -> PregelNode: + ... @overload @classmethod @@ -126,7 +127,8 @@ class Channel: *, key: None = None, tags: Optional[list[str]] = None, - ) -> PregelNode: ... + ) -> PregelNode: + ... @classmethod def subscribe_to( @@ -1210,7 +1212,8 @@ def _prepare_next_tasks( processes: Mapping[str, PregelNode], channels: Mapping[str, BaseChannel], for_execution: Literal[False], -) -> tuple[Checkpoint, list[PregelTaskDescription]]: ... +) -> tuple[Checkpoint, list[PregelTaskDescription]]: + ... @overload @@ -1219,7 +1222,8 @@ def _prepare_next_tasks( processes: Mapping[str, PregelNode], channels: Mapping[str, BaseChannel], for_execution: Literal[True], -) -> tuple[Checkpoint, list[PregelExecutableTask]]: ... +) -> tuple[Checkpoint, list[PregelExecutableTask]]: + ... def _prepare_next_tasks( diff --git a/langgraph/serde/base.py b/langgraph/serde/base.py index 5f1250d1e..2fbb0ab71 100644 --- a/langgraph/serde/base.py +++ b/langgraph/serde/base.py @@ -10,6 +10,8 @@ class SerializerProtocol(Protocol): Valid implementations include the `pickle`, `json` and `orjson` modules. """ - def dumps(self, obj: Any) -> bytes: ... + def dumps(self, obj: Any) -> bytes: + ... - def loads(self, data: bytes) -> Any: ... + def loads(self, data: bytes) -> Any: + ... diff --git a/langgraph/version.py b/langgraph/version.py index 3368893c0..ac7aeef6f 100644 --- a/langgraph/version.py +++ b/langgraph/version.py @@ -1,5 +1,4 @@ """Main entrypoint into package.""" - from importlib import metadata try: From aa48a3be3d3a467d374cb3c4cf2ad3856779e283 Mon Sep 17 00:00:00 2001 From: Harrison Chase Date: Mon, 29 Apr 2024 16:37:16 -0700 Subject: [PATCH 3/4] cr --- langgraph/prebuilt/chat_agent_executor.py | 15 ++++++++++++++- langgraph/version.py | 1 + 2 files changed, 15 insertions(+), 1 deletion(-) diff --git a/langgraph/prebuilt/chat_agent_executor.py b/langgraph/prebuilt/chat_agent_executor.py index 03d90c890..453e06438 100644 --- a/langgraph/prebuilt/chat_agent_executor.py +++ b/langgraph/prebuilt/chat_agent_executor.py @@ -2,7 +2,7 @@ import json from typing import Annotated, Optional, Sequence, TypedDict, Union from langchain_core.language_models import LanguageModelLike -from langchain_core.messages import BaseMessage, FunctionMessage +from langchain_core.messages import BaseMessage, FunctionMessage, SystemMessage from langchain_core.runnables import RunnableLambda from langchain_core.tools import BaseTool from langchain_core.utils.function_calling import convert_to_openai_function @@ -137,6 +137,7 @@ def create_function_calling_executor( def create_tool_calling_executor( model: LanguageModelLike, tools: Union[ToolExecutor, Sequence[BaseTool]], + system_message: Optional[Union[str, SystemMessage]] = None, checkpointer: Optional[BaseCheckpointSaver] = None, interrupt_before: Optional[Sequence[str]] = None, interrupt_after: Optional[Sequence[str]] = None, @@ -147,6 +148,12 @@ def create_tool_calling_executor( Args: model (LanguageModelLike): The chat model that supports OpenAI tool calling. tools (Union[ToolExecutor, Sequence[BaseTool]]): A list of tools or a ToolExecutor instance. + system_message: (Optional[Union[str, SystemMessage]]): An optional system message to pass in + to the model. Is appended at the start of the messages. + checkpointer (Optional[BaseCheckpointSaver]): An optional checkpoint saver object. + interrupt_before (Optional[Sequence[str]]): An optional list of node names to interrupt before. + interrupt_after (Optional[Sequence[str]]): An optional list of node names to interrupt after. + debug (bool): A flag indicating whether to enable debug mode. Returns: Runnable: A compiled LangChain runnable that can be used for chat interactions. @@ -188,6 +195,12 @@ def create_tool_calling_executor( # Define the function that calls the model def call_model(state: AgentState): messages = state["messages"] + if system_message is not None: + if isinstance(system_message, str): + _system_message: BaseMessage = SystemMessage(content=system_message) + else: + _system_message = system_message + messages = [_system_message] + list(messages) response = model.invoke(messages) # We return a list, because this will get added to the existing list return {"messages": [response]} diff --git a/langgraph/version.py b/langgraph/version.py index ac7aeef6f..3368893c0 100644 --- a/langgraph/version.py +++ b/langgraph/version.py @@ -1,4 +1,5 @@ """Main entrypoint into package.""" + from importlib import metadata try: From 97f3d66c0e6642e1b31ce0d2d3cfa7b1d78c8a97 Mon Sep 17 00:00:00 2001 From: Harrison Chase Date: Wed, 1 May 2024 08:17:21 -0700 Subject: [PATCH 4/4] cr --- langgraph/prebuilt/chat_agent_executor.py | 40 +++++++--- tests/test_prebuilt.py | 97 +++++++++++++++++++++++ 2 files changed, 124 insertions(+), 13 deletions(-) create mode 100644 tests/test_prebuilt.py diff --git a/langgraph/prebuilt/chat_agent_executor.py b/langgraph/prebuilt/chat_agent_executor.py index 453e06438..4632e3488 100644 --- a/langgraph/prebuilt/chat_agent_executor.py +++ b/langgraph/prebuilt/chat_agent_executor.py @@ -1,9 +1,9 @@ import json -from typing import Annotated, Optional, Sequence, TypedDict, Union +from typing import Annotated, Callable, Optional, Sequence, TypedDict, Union from langchain_core.language_models import LanguageModelLike from langchain_core.messages import BaseMessage, FunctionMessage, SystemMessage -from langchain_core.runnables import RunnableLambda +from langchain_core.runnables import Runnable, RunnableLambda from langchain_core.tools import BaseTool from langchain_core.utils.function_calling import convert_to_openai_function @@ -137,7 +137,7 @@ def create_function_calling_executor( def create_tool_calling_executor( model: LanguageModelLike, tools: Union[ToolExecutor, Sequence[BaseTool]], - system_message: Optional[Union[str, SystemMessage]] = None, + messages_modifier: Optional[Union[SystemMessage, str, Callable, Runnable]] = None, checkpointer: Optional[BaseCheckpointSaver] = None, interrupt_before: Optional[Sequence[str]] = None, interrupt_after: Optional[Sequence[str]] = None, @@ -148,8 +148,13 @@ def create_tool_calling_executor( Args: model (LanguageModelLike): The chat model that supports OpenAI tool calling. tools (Union[ToolExecutor, Sequence[BaseTool]]): A list of tools or a ToolExecutor instance. - system_message: (Optional[Union[str, SystemMessage]]): An optional system message to pass in - to the model. Is appended at the start of the messages. + messages_modifier: (Optional[Union[SystemMessage, str, Callable, Runnable]]): An optional + messages modifier. This applies to messages BEFORE they are passed into the LLM. + Can take a few different forms: + - SystemMessage: this is added to the beginning of the list of messages. + - str: This is converted to a SystemMessage and added to the beginning of the list of messages. + - Callable: This function should take in a list of messages and the output is then passed to the language model. + - Runnable: This runnable should take in a list of messages and the output is then passed to the language model. checkpointer (Optional[BaseCheckpointSaver]): An optional checkpoint saver object. interrupt_before (Optional[Sequence[str]]): An optional list of node names to interrupt before. interrupt_after (Optional[Sequence[str]]): An optional list of node names to interrupt after. @@ -192,22 +197,31 @@ def create_tool_calling_executor( else: return "continue" + # Add the message modifier, if exists + if messages_modifier is None: + model_runnable = model + elif isinstance(messages_modifier, str): + _system_message: BaseMessage = SystemMessage(content=messages_modifier) + model_runnable = (lambda messages: [_system_message] + messages) | model + elif isinstance(messages_modifier, SystemMessage): + model_runnable = (lambda messages: [messages_modifier] + messages) | model + elif isinstance(messages_modifier, (Callable, Runnable)): + model_runnable = messages_modifier | model + else: + raise ValueError( + f"Got unexpected type for `messages_modifier`: {type(messages_modifier)}" + ) + # Define the function that calls the model def call_model(state: AgentState): messages = state["messages"] - if system_message is not None: - if isinstance(system_message, str): - _system_message: BaseMessage = SystemMessage(content=system_message) - else: - _system_message = system_message - messages = [_system_message] + list(messages) - response = model.invoke(messages) + response = model_runnable.invoke(messages) # We return a list, because this will get added to the existing list return {"messages": [response]} async def acall_model(state: AgentState): messages = state["messages"] - response = await model.ainvoke(messages) + response = await model_runnable.ainvoke(messages) # We return a list, because this will get added to the existing list return {"messages": [response]} diff --git a/tests/test_prebuilt.py b/tests/test_prebuilt.py new file mode 100644 index 000000000..57a1d1159 --- /dev/null +++ b/tests/test_prebuilt.py @@ -0,0 +1,97 @@ +from typing import Any, Callable, Dict, List, Optional, Sequence, Type, Union + +from langchain_core.callbacks import ( + CallbackManagerForLLMRun, +) +from langchain_core.language_models import ( + BaseChatModel, + LanguageModelInput, +) +from langchain_core.messages import AIMessage, BaseMessage, HumanMessage, SystemMessage +from langchain_core.outputs import ChatGeneration, ChatResult +from langchain_core.pydantic_v1 import BaseModel +from langchain_core.runnables import Runnable, RunnableLambda +from langchain_core.tools import BaseTool + +from langgraph.prebuilt.chat_agent_executor import create_tool_calling_executor + + +class FakeToolCallingModel(BaseChatModel): + def _generate( + self, + messages: List[BaseMessage], + stop: Optional[List[str]] = None, + run_manager: Optional[CallbackManagerForLLMRun] = None, + **kwargs: Any, + ) -> ChatResult: + """Top Level call""" + messages_string = "-".join([m.content for m in messages]) + message = AIMessage(content=messages_string, id="0") + return ChatResult(generations=[ChatGeneration(message=message)]) + + @property + def _llm_type(self) -> str: + return "fake-tool-call-model" + + def bind_tools( + self, + tools: Sequence[Union[Dict[str, Any], Type[BaseModel], Callable, BaseTool]], + **kwargs: Any, + ) -> Runnable[LanguageModelInput, BaseMessage]: + if len(tools) > 0: + raise ValueError("Not supported yet!") + return self + + +def test_no_modifier(): + model = FakeToolCallingModel() + agent = create_tool_calling_executor(model, []) + inputs = [HumanMessage("hi?")] + response = agent.invoke({"messages": inputs}) + expected_response = {"messages": inputs + [AIMessage(content="hi?", id="0")]} + assert response == expected_response + + +def test_system_message_modifier(): + model = FakeToolCallingModel() + messages_modifier = SystemMessage(content="Foo") + agent = create_tool_calling_executor(model, [], messages_modifier=messages_modifier) + inputs = [HumanMessage("hi?")] + response = agent.invoke({"messages": inputs}) + expected_response = {"messages": inputs + [AIMessage(content="Foo-hi?", id="0")]} + assert response == expected_response + + +def test_system_message_string_modifier(): + model = FakeToolCallingModel() + messages_modifier = "Foo" + agent = create_tool_calling_executor(model, [], messages_modifier=messages_modifier) + inputs = [HumanMessage("hi?")] + response = agent.invoke({"messages": inputs}) + expected_response = {"messages": inputs + [AIMessage(content="Foo-hi?", id="0")]} + assert response == expected_response + + +def test_callable_modifier(): + model = FakeToolCallingModel() + + def messages_modifier(messages): + return [HumanMessage(content="Bar")] + + agent = create_tool_calling_executor(model, [], messages_modifier=messages_modifier) + inputs = [HumanMessage("hi?")] + response = agent.invoke({"messages": inputs}) + expected_response = {"messages": inputs + [AIMessage(content="Bar", id="0")]} + assert response == expected_response + + +def test_runnable_modifier(): + model = FakeToolCallingModel() + + messages_modifier = RunnableLambda(lambda x: [HumanMessage(content="Baz")]) + + agent = create_tool_calling_executor(model, [], messages_modifier=messages_modifier) + inputs = [HumanMessage("hi?")] + response = agent.invoke({"messages": inputs}) + expected_response = {"messages": inputs + [AIMessage(content="Baz", id="0")]} + assert response == expected_response